make_zero! can sometimes fail on closures?

[ChrisRackauckas]

It's the error noticed in SciML/SciMLSensitivity.jl#1067

  LoadError: MethodError: no method matching iterate(::SciMLSensitivity.var"#136#140"{SciMLBase.ODEFunction{true, true, SciMLBase.ODEFunction{true, SciMLBase.AutoSpecialize, Main.var"##Enzyme Closures#265".var"#dfunc#3"{Main.var"##Enzyme Closures#265".NN{Float64}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}})
  
  Closest candidates are:
    iterate(::DataStructures.TrieIterator)
     @ DataStructures ~/.julia/packages/DataStructures/95DJa/src/trie.jl:112
    iterate(::DataStructures.TrieIterator, ::Any)
     @ DataStructures ~/.julia/packages/DataStructures/95DJa/src/trie.jl:112
    iterate(::Tables.DictRowTable)
     @ Tables ~/.julia/packages/Tables/NSGZI/src/dicts.jl:122
    ...
  
  Stacktrace:
    [1] in(x::Base.IdSet{Any}, itr::Function)
      @ Base ./operators.jl:1292
    [2] make_zero!
      @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:1495 [inlined]
    [3] make_zero!
      @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:1492 [inlined]

Cutting off the SciML parts.

For some reason this function fails Enzyme.Compiler.guaranteed_const_nongen(typeof(_pf), nothing) false and then in the next iterate it errors. It seems like make_zero! is not robust to active values in closures?

For some reason my attempts at a simple MWE have failed. Maybe @wsmoses knows what's going on, it seems guaranteed_const_nongen can be rather particular based on the environment or something.

[wsmoses]

Can you paste the whole log and inputs.

im probably going to need a reproducer to fix

[wsmoses]

Im not sure it’s the case here, but one comment I’ll add — make_zero! Can only set the mutable data eg you can’t do make_zero!(1.0)

That doesn’t mean the datastructure you pass in needs to be mutable, you could do make_zero!(([1.],[2.])) even though the data type is a tuple, there is mutable data in the arrays in the tuple. This would fail however for a tuple of floats.

[ChrisRackauckas]

using Enzyme
struct Dense
    n_inp::Int
    b::Vector{Float64}
end

function Dense(n)
    Dense(n, rand(n))
end

nn = Dense(4)
Enzyme.make_zero!(nn)

ERROR: MethodError: no method matching iterate(::Dense)

Closest candidates are:
  iterate(::LLVM.FunctionParameterSet)
   @ LLVM ~/.julia/packages/LLVM/6cDbl/src/core/function.jl:103
  iterate(::LLVM.FunctionParameterSet, ::Any)
   @ LLVM ~/.julia/packages/LLVM/6cDbl/src/core/function.jl:103
  iterate(::Core.Compiler.InstructionStream, ::Int64)
   @ Base show.jl:2778
  ...

Stacktrace:
 [1] in(x::Base.IdSet{Any}, itr::Dense)
   @ Base ./operators.jl:1292
 [2] make_zero!
   @ ~/.julia/packages/Enzyme/qd8AI/src/compiler.jl:1576 [inlined]
 [3] make_zero!(prev::Dense)
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/qd8AI/src/compiler.jl:1573
 [4] top-level scope
   @ ~/.julia/dev/SciMLSensitivity/test/enzyme_closure.jl:12

[wsmoses]

Okay this should be fixable, will take a look at shortly

[wsmoses]

Fixe don main