besseli (and friends?)

[willtebbutt]

using SpecialFunctions
using Zygote

Zygote.gradient(besseli, 4, 0.3)

yields

ERROR: not implemented
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] (::SpecialFunctions.var"#27#30")()
    @ SpecialFunctions ~/.julia/packages/SpecialFunctions/mFAQ4/src/chainrules.jl:41
  [3] (::ChainRulesCore.Thunk{SpecialFunctions.var"#27#30"})()
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/D0go7/src/differentials/thunks.jl:98
  [4] unthunk
    @ ~/.julia/packages/ChainRulesCore/D0go7/src/differentials/thunks.jl:99 [inlined]
  [5] (::ChainRulesCore.var"#11#12"{ChainRulesCore.Thunk{SpecialFunctions.var"#27#30"}})()
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/D0go7/src/differentials/thunks.jl:40
  [6] (::ChainRulesCore.Thunk{ChainRulesCore.var"#11#12"{ChainRulesCore.Thunk{SpecialFunctions.var"#27#30"}}})()
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/D0go7/src/differentials/thunks.jl:98
  [7] unthunk
    @ ~/.julia/packages/ChainRulesCore/D0go7/src/differentials/thunks.jl:99 [inlined]
  [8] *(a::ChainRulesCore.Thunk{ChainRulesCore.var"#11#12"{ChainRulesCore.Thunk{SpecialFunctions.var"#27#30"}}}, b::Float64)
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/D0go7/src/differential_arithmetic.jl:111
  [9] besseli_pullback
    @ ~/.julia/packages/SpecialFunctions/mFAQ4/src/chainrules.jl:38 [inlined]
 [10] ZBack
    @ ~/.julia/packages/Zygote/RxTZu/src/compiler/chainrules.jl:77 [inlined]
 [11] (::Zygote.var"#41#42"{Zygote.ZBack{SpecialFunctions.var"#besseli_pullback#29"{Int64, Float64}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/RxTZu/src/compiler/interface.jl:41
 [12] gradient(::Function, ::Int64, ::Vararg{Any, N} where N)
    @ Zygote ~/.julia/packages/Zygote/RxTZu/src/compiler/interface.jl:59
 [13] top-level scope
    @ REPL[4]:1

Perhaps @oxinabox or @mzgubic could confirm whether the following explanation of the problem is accurate:

Zygote doesn't do thunks, so when converting between Zygote and ChainRules types we strip out any thunks by unthunk`ing them.

At the same time, there are some rules in ChainRules where we don't have cotangents implemented w.r.t. certain arguments. For such arguments, we return a thunk, which throws the above error if it's ever actually used. besseli is one such function.

The thinking is that if you never actually need the gradient w.r.t. the argument whose gradient we don't know, we can avoid ever having to worry about not knowing how to compute it. Alas, unthunking to get a type that Zygote can handle means we always evaluate the thing.

Any thoughts on good resolutions?

This is going to cause problems for some TemporalGPs.jl work in the next couple of days. I can add a workaround for now, but that's obviously not much use for anyone else who runs into this.

edit: cross-referenced relevant TemporalGPs issue.

[mzgubic]

The description sounds right to me. Probably the best way to solve this is to support ChainRules types internally in Zygote, as described in #603. I guess an alternative is to add an @adjoint for the function, but that moves us in the opposite direction that we want to go (i.e. more rrules, less adjoints)

[DhairyaLGandhi]

Hey I like adjoints :p

[oxinabox]

@devmotion proposed a most short term solution yesterday.
Introduce a type specifically for NotImplemented.
JuliaDiff/ChainRulesCore.jl#334

we probably want that even once we do #603

[devmotion]

The Zygote issue is also discussed in #873.

[willtebbutt]

Closing in favour of #873