# ChainRulesCore: Custom adjoint ignored in certain cases

**URL:** <https://discourse.julialang.org/t/chainrulescore-custom-adjoint-ignored-in-certain-cases/67814>\
**Category:** Machine Learning\
**Tags:** adjoint, autodiff, chainrulescore\
**Created:** [September 7, 2021, 1:01pm UTC](https://discourse.julialang.org/t/chainrulescore-custom-adjoint-ignored-in-certain-cases/67814 "2021-09-07T13:01:35Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Rasmus\_Hoier](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rasmus_hoier/32/24036_2.png) [@Rasmus\_Hoier](https://discourse.julialang.org/u/Rasmus_Hoier)\
**Post date:** [September 7, 2021, 1:01pm UTC](https://discourse.julialang.org/t/chainrulescore-custom-adjoint-ignored-in-certain-cases/67814/1 "2021-09-07T13:01:35Z")

</div>

Hi,  
A colleague of mine ran into the issue that custom adjoints get ignored in certain cases in recent versions of ChainRulesCore. Here is a minimal working example:

```julia
using Flux, ChainRulesCore

       my_fun(x) = x*x

       function ChainRulesCore.rrule(::typeof(my_fun), x)
           z = my_fun(x)
           function pullback(Δy)
                  @show "Calling custom pullback"
                  return NoTangent(), zero(x)
           end
           return z, pullback
       end

       @show gradient(my_fun, 1.0) # calls custom rrule
       @show gradient(x -> sum(my_fun.(x)), [1.0]) # does not
       @show gradient(x -> sum(map(my_fun, x)), [1.0]) # calls custom rrule

```

When running this example with ChainRulesCore version 0.10.13 or older the custom gradient is used in all three cases. But when running the example using ChainRulesCore version 1.0.1 or newer the custom gradient is ignored in the second case. I ran into problems when trying to install ChainRulesCore version 1.0.0, so I could not test this version.

Is this a deliberate design choice or a bug?

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [September 7, 2021, 1:18pm UTC](https://discourse.julialang.org/t/chainrulescore-custom-adjoint-ignored-in-certain-cases/67814/2 "2021-09-07T13:18:38Z")

</div>

It’s a performance trick.

When broadcasting, it checks `Base.issingletontype(typeof(my_fun))`, and concludes that the function is pure enough that it can use ForwardDiff within the broadcast, rather than Zygote’s own differentiation (which calls ChainRules). This is often 100x faster, because Zygote’s differentiation tends to break type stability.

You can write a method `my_fun(x::ForwardDiff.Dual)` to customise this. Or you can define a rule for the broadcasting, Zygote has many rules which look like this:

```julia
@adjoint function broadcasted(::typeof(tanh), x::Numeric)
  y = tanh.(x)
  y, ȳ -> (nothing, ȳ .* conj.(1 .- y.^2))
end

```

---

<div class="post-metadata">

**Author:** ![Rasmus\_Hoier](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rasmus_hoier/32/24036_2.png) [@Rasmus\_Hoier](https://discourse.julialang.org/u/Rasmus_Hoier)\
**Post date:** [September 7, 2021, 2:21pm UTC](https://discourse.julialang.org/t/chainrulescore-custom-adjoint-ignored-in-certain-cases/67814/3 "2021-09-07T14:21:10Z")

</div>

Thanks!  
Pretty neat that the call to gradient checks if forward mode will be beneficial! 🙂 I guess we had implicitly made the assumption that the gradient function would always use reverse mode, which is not really justified.

BTW: I was a bit confused about why use ForwardDiff instead of ChainRulesCore.frule. Found an answer to this [here](https://discourse.julialang.org/t/chainrulescore-and-forwarddiff/61705).

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [September 7, 2021, 2:32pm UTC](https://discourse.julialang.org/t/chainrulescore-custom-adjoint-ignored-in-certain-cases/67814/4 "2021-09-07T14:32:10Z")

</div>

Right, so ForwardDiff doesn’t use ChainRules (it’s old & stable, and maybe wouldn’t fit well). And the broadcasting needs to work even if the function is something you’re written, or some composition of functions, not just elementary functions. I suppose it could check for the existence of a direct `frule` first but someone would have to write that.

It is bit of a pain that you potentially have to write 3 different rules for the same function, sorry about the surprise! I think we may still be some way from the optimal story here, although not many more factors of 100 surely.