# How to deal with Zygote sometimes "pirating" its own adjoints with worse ones?

**URL:** <https://discourse.julialang.org/t/how-to-deal-with-zygote-sometimes-pirating-its-own-adjoints-with-worse-ones/32579>\
**Category:** General Usage\
**Created:** [December 22, 2019, 3:02pm UTC](https://discourse.julialang.org/t/how-to-deal-with-zygote-sometimes-pirating-its-own-adjoints-with-worse-ones/32579 "2019-12-22T15:02:56Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![marius311](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marius311/32/3953_2.png) [@marius311](https://discourse.julialang.org/u/marius311)\
**Post date:** [December 22, 2019, 3:02pm UTC](https://discourse.julialang.org/t/how-to-deal-with-zygote-sometimes-pirating-its-own-adjoints-with-worse-ones/32579/1 "2019-12-22T15:02:56Z")

</div>

The situation I’m in is that I have some custom arrays and some custom implementation of multiplication, e.g.

```julia
function (A::CustomMatrix * b::CustomVector)
    some_specialized_code()
end

```

Now if I use Zygote to derivative through code which does `A*b` somewhere, what will happen is that Zygote uses the adjoint defined [here](https://github.com/FluxML/Zygote.jl/blob/727ea3c70eb276fdb3b82e0b2c802b8c52f4724a/src/lib/array.jl#L244-L246), which is:

```julia
@adjoint function(A::AbstractMatrix * x::AbstractVector)
  return A * x, Δ::AbstractVector->(Δ * x', A' * Δ)
end

```

But in this particular case, the resulting adjoint is extremely inefficient as compared to if Zygote had simply taken a gradient through my custom implementation in `some_specialized_code()`.

My question is whether there is a way to keep Zygote from “pirating” its own adjoint like this, or whether there is an easy way to write a custom adjoint that forwards the adjoint to `some_specialized_code()` instead? Thanks.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 22, 2019, 3:23pm UTC](https://discourse.julialang.org/t/how-to-deal-with-zygote-sometimes-pirating-its-own-adjoints-with-worse-ones/32579/2 "2019-12-22T15:23:16Z")

</div>

```julia
function (A::CustomMatrix * b::CustomVector)
    some_specialized_code(A,b)
end

@adjoint function (A::CustomMatrix * b::CustomVector)
    Zygote.pullback(some_specialized_code,A,b)
end

```

---

<div class="post-metadata">

**Author:** ![marius311](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marius311/32/3953_2.png) [@marius311](https://discourse.julialang.org/u/marius311)\
**Post date:** [December 22, 2019, 3:35pm UTC](https://discourse.julialang.org/t/how-to-deal-with-zygote-sometimes-pirating-its-own-adjoints-with-worse-ones/32579/3 "2019-12-22T15:35:17Z")

</div>

Perfect, thanks!

---

<div class="post-metadata">

**Author:** ![marius311](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marius311/32/3953_2.png) [@marius311](https://discourse.julialang.org/u/marius311)\
**Post date:** [December 24, 2019, 3:14pm UTC](https://discourse.julialang.org/t/how-to-deal-with-zygote-sometimes-pirating-its-own-adjoints-with-worse-ones/32579/4 "2019-12-24T15:14:33Z")

</div>

Followup question, but is there any even more automatic way to do it that doesn’t refer to `some_specialized_code()`? I ask because in some cases, I don’t actually have access to this function as its just a function in `Base`. This happens e.g. for `b' * A`. If Zygote derived through `Base`’s definition which is `(A'*b)'`, things would work fine since `A' isa CustomMatrix` so it’d fall to the above adjoint definition suggested by Chris. Instead though its using Zygote’s own definition for adjoint of `(::AbstractMatrix * ::AbstractMatrix)` which is now inefficient.

I guess ideally I need something like:

```julia
@adjoint function (b_adj::Adjoint{<:Any, CustomVector} * A::CustomMatrix)
    # teach Zygote that for this call, explicitly derive 
    # through the code that `b_adj * A` actually calls, rather 
    # than using the rule for adjoint of `b_adj * A`
end

```

Is anything like this possible? Or maybe there’s a better design that side-steps the need for this? Thanks.