Zygote can't differentiate IndexLens

[penelopeysm]

Run the following:

using DynamicPPL, Distributions, ADTypes
using LogDensityProblems: logdensity_and_gradient
import Zygote

@model function observe_index(x=[1.5, 2.0, 2.5])
    a ~ Normal()
    for i in eachindex(x)
        x[i] ~ Normal(a)
    end
end

function grad(adtype)
    logdensity_and_gradient(LogDensityFunction(observe_index(); adtype=adtype), [2.0])
end

grad(AutoZygote())

Result:

ERROR: Need an adjoint for constructor Accessors.IndexLens{Tuple{Int64}}. Gradient is of type ChainRulesCore.ZeroTangent
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] (::Zygote.Jnew{Accessors.IndexLens{Tuple{Int64}}, Nothing, false})(Δ::ChainRulesCore.ZeroTangent)
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/lib/lib.jl:324
  [3] (::Zygote.var"#2230#back#327"{Zygote.Jnew{Accessors.IndexLens{…}, Nothing, false}})(Δ::ChainRulesCore.ZeroTangent)
    @ Zygote ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72
  [4] IndexLens
    @ ~/.julia/packages/Accessors/8Y4uq/src/optics.jl:402 [inlined]
  [5] (::Zygote.Pullback{Tuple{Type{…}, Tuple{…}}, Tuple{Zygote.var"#2230#back#327"{…}}})(Δ::ChainRulesCore.ZeroTangent)
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
  [6] observe_index
    @ ~/.julia/packages/DynamicPPL/l7W1f/src/compiler.jl:574 [inlined]
  [7] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Tuple{Nothing, @NamedTuple{…}})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
  [8] #305
    @ ~/.julia/packages/Zygote/HdT4O/src/lib/lib.jl:202 [inlined]
  [9] (::Zygote.var"#2189#back#307"{Zygote.var"#305#306"{…}})(Δ::Tuple{Nothing, @NamedTuple{…}})
    @ Zygote ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72
 [10] _evaluate!!
    @ ~/.julia/packages/DynamicPPL/l7W1f/src/model.jl:936 [inlined]
 [11] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Tuple{Nothing, @NamedTuple{…}})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
 [12] evaluate_threadunsafe!!
    @ ~/.julia/packages/DynamicPPL/l7W1f/src/model.jl:909 [inlined]
 [13] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{Nothing, @NamedTuple{…}})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
 [14] evaluate!!
    @ ~/.julia/packages/DynamicPPL/l7W1f/src/model.jl:857 [inlined]
 [15] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Tuple{Nothing, @NamedTuple{…}})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
 [16] logdensity_at
    @ ~/.julia/packages/DynamicPPL/l7W1f/src/logdensityfunction.jl:184 [inlined]
 [17] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
 [18] (::Zygote.var"#88#89"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface.jl:97
 [19] withgradient(::Function, ::Vector{Float64}, ::Vararg{Any})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface.jl:0
 [20] value_and_gradient(::Function, ::DifferentiationInterface.NoGradientPrep{…}, ::AutoZygote, ::Vector{…}, ::DifferentiationInterface.Constant{…}, ::DifferentiationInterface.Constant{…}, ::DifferentiationInterface.Constant{…})
    @ DifferentiationInterfaceZygoteExt ~/.julia/packages/DifferentiationInterface/7eD1K/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl:115
 [21] logdensity_and_gradient(f::LogDensityFunction{…}, x::Vector{…})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/l7W1f/src/logdensityfunction.jl:215
 [22] grad(adtype::AutoZygote)
    @ Main ./REPL[9]:2
 [23] top-level scope
    @ REPL[10]:1
Some type information was truncated. Use `show(err)` to see complete types.

[penelopeysm]

Oh, that's because it's mutating arrays. Of course.

[penelopeysm]

Wait, it's not mutating arrays, right? :/ Especially not for observe. For assume it would be mutating.