InplaceableThunk is missing implementation of keys and eachindex

[MilesCranmer]

I got this stack trace when trying to upgrade Zygote.jl in DynamicExpressions.jl in SymbolicML/DynamicExpressions.jl#119. InplaceableThunk seems to be indexable, but may be missing the implementation of keys and eachindex. (Edit: and also length)

I can work around it but wanted to point out the missing interface.

This is SHA 40a3130dead045f2428909380e782ca10e7851db in github.com/SymbolicML/DynamicExpressions.jl in case you would like to reproduce it.

julia> using DynamicExpressions

julia> using Zygote: Zygote

julia> using DifferentiationInterface: AutoZygote, gradient

julia> ex = @parse_expression(x1 + 1.5, binary_operators = [+], variable_names = ["x1"])
x1 + 1.5

julia> d_ex = gradient(AutoZygote(), ex) do ex
           sum(ex(ones(1, 5)))
       end
ERROR: MethodError: no method matching keys(::ChainRulesCore.InplaceableThunk{ChainRulesCore.Thunk{…}, ChainRules.var"#718#721"{…}}, ::Base.OneTo{Int64})
The function `keys` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  keys(::Zygote.Buffer, ::Any...; kwargs...)
   @ Zygote ~/.julia/packages/MacroTools/Ar0jT/src/examples/forward.jl:17
  keys(::IndexStyle, ::AbstractArray, AbstractArray...)
   @ Base abstractarray.jl:404
  keys(::DataStructures.MultiDict, ::Any...)
   @ DataStructures ~/.julia/packages/DataStructures/IrAJn/src/delegate.jl:21
  ...

Stacktrace:
  [1] eachindex(::ChainRulesCore.InplaceableThunk{ChainRulesCore.Thunk{…}, ChainRules.var"#718#721"{…}}, ::Base.OneTo{Int64})
    @ Base ./abstractarray.jl:318
  [2] (::DynamicExpressions.ChainRulesModule.EvalPullback{…})(::ChainRulesCore.Tangent{…})
    @ DynamicExpressions.ChainRulesModule ~/PermaDocuments/SymbolicRegressionMonorepo/DynamicExpressions.jl/src/ChainRules.jl:70
  [3] (::Zygote.ZBack{…})(dy::Tuple{…})
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/chainrules.jl:222
  [4] #_#1
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/DynamicExpressions.jl/src/EvaluationHelpers.jl:29 [inlined]
  [5] (::Zygote.Pullback{…})(Δ::ChainRulesCore.InplaceableThunk{…})
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/interface2.jl:0
  [6] AbstractExpressionNode
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/DynamicExpressions.jl/src/EvaluationHelpers.jl:28 [inlined]
  [7] (::Zygote.Pullback{…})(Δ::ChainRulesCore.InplaceableThunk{…})
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/interface2.jl:0
  [8] #_#23
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/DynamicExpressions.jl/src/Expression.jl:513 [inlined]
  [9] (::Zygote.Pullback{…})(Δ::ChainRulesCore.InplaceableThunk{…})
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/interface2.jl:0
 [10] AbstractExpression
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/DynamicExpressions.jl/src/Expression.jl:509 [inlined]
 [11] (::Zygote.Pullback{…})(Δ::ChainRulesCore.InplaceableThunk{…})
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/interface2.jl:0
 [12] AbstractExpression
    @ ~/PermaDocuments/SymbolicRegressionMonorepo/DynamicExpressions.jl/src/Expression.jl:509 [inlined]
 [13] #550
    @ ./REPL[13]:2 [inlined]
 [14] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/interface2.jl:0
 [15] (::Zygote.var"#88#89"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/interface.jl:97
 [16] gradient(f::Function, args::Expression{Float64, Node{…}, @NamedTuple{…}})
    @ Zygote ~/.julia/packages/Zygote/wfLOG/src/compiler/interface.jl:154
 [17] gradient
    @ ~/.julia/packages/DifferentiationInterface/zJHX8/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl:123 [inlined]
 [18] gradient(::var"#550#551", ::AutoZygote, ::Expression{Float64, Node{…}, @NamedTuple{…}})
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/zJHX8/src/first_order/gradient.jl:63
 [19] top-level scope
    @ REPL[13]:1
Some type information was truncated. Use `show(err)` to see complete types.

It could be a different issue but I found it surprising I could index this object despite eachindex not being defined.

[mcabbott]

Some methods automatically un-thunk, for convenience. But making getindex do so sounds like a mistake, because typically you will do this more than once, and it is wasteful to do the deferred calculation many times:

julia> dx = ChainRules.@thunk [1,2,3] .+ something(sleep(1), 4)
Thunk(var"#13#14"())

julia> @time dx[1]
  1.026216 seconds (100.73 k allocations: 5.205 MiB, 2.34% compilation time)
5

julia> @time dx[2]
  1.002163 seconds (8 allocations: 272 bytes)
6

julia> eachindex(dx)
ERROR: MethodError: no method matching keys(::ChainRulesCore.Thunk{var"#13#14"})

So my suggestion would be not to define eachindex. And instead to call unthunk(dY) once, then use the result in the loop, in whatever rule is causing this. Perhaps here:

https://github.com/SymbolicML/DynamicExpressions.jl/blob/master/src/ChainRules.jl#L55C29-L75

(This wasn't an issue before Zygote 0.7, as all Thunks were immediately un-thunked.)

[MilesCranmer]

(This wasn't an issue before Zygote 0.7, as all Thunks were immediately un-thunked.)

Ah, this must be the missing piece. Thanks!

[MilesCranmer]

Fix was just

- function (e::EvalPullback)((dY, _))
+ function (e::EvalPullback)((thunked_dY, _))
+     dY = unthunk(thunked_dY)