Autodiff code generation fails when caching values

[avik-pal]

using Reactant, Enzyme

mutable struct StateReturn
    st::Any
end

function simple_ad_fail(x, stret)
    loss = sum(x)
    stret.st = x .+ 1
    return loss
end

x = rand(10)
x_ra = x |> Reactant.to_rarray

Enzyme.gradient(Reverse, simple_ad_fail, x, Const(StateReturn(nothing)))

@code_hlo Enzyme.gradient(Reverse, simple_ad_fail, x_ra, Const(StateReturn(nothing)))

ERROR: AssertionError: No path found for TracedRArray{Float64,1N}(((1,),), size=(10,))
Stacktrace:
  [1] get_argidx(x::Reactant.TracedRArray{Float64, 1})
    @ Reactant.TracedUtils /mnt/software/lux/Reactant.jl/src/TracedUtils.jl:341
  [2] overload_autodiff(::ReverseMode{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…})
    @ Reactant /mnt/software/lux/Reactant.jl/src/Interpreter.jl:306
  [3] autodiff(::ReverseMode{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…})
    @ Reactant /mnt/software/lux/Reactant.jl/src/Overlay.jl:32
  [4] autodiff
    @ /mnt/.julia/packages/Enzyme/ydGh2/src/Enzyme.jl:524 [inlined]
  [5] macro expansion
    @ /mnt/.julia/packages/Enzyme/ydGh2/src/sugar.jl:270 [inlined]
  [6] gradient
    @ /mnt/.julia/packages/Enzyme/ydGh2/src/sugar.jl:258 [inlined]
  [7] gradient(rm::ReverseMode{…}, f::typeof(simple_ad_fail), x::Reactant.TracedRArray{…}, args::Tuple{…})
    @ Reactant ./<missing>:0
  [8] GenericMemory
    @ ./boot.jl:516 [inlined]
  [9] IdDict
    @ ./iddict.jl:31 [inlined]
 [10] IdDict
    @ ./iddict.jl:49 [inlined]
 [11] make_zero (repeats 2 times)
    @ /mnt/.julia/packages/EnzymeCore/15Zff/src/EnzymeCore.jl:529 [inlined]
 [12] macro expansion
    @ /mnt/.julia/packages/Enzyme/ydGh2/src/sugar.jl:316 [inlined]
 [13] gradient
    @ /mnt/.julia/packages/Enzyme/ydGh2/src/sugar.jl:258 [inlined]
 [14] (::Reactant.TracedUtils.var"#8#18"{Bool, Bool, typeof(Enzyme.gradient), Tuple{…}, Vector{…}, Tuple{…}})()
    @ Reactant.TracedUtils /mnt/software/lux/Reactant.jl/src/TracedUtils.jl:210
 [15] block!(f::Reactant.TracedUtils.var"#8#18"{…}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201
 [16] make_mlir_fn(f::Function, args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, no_args_in_result::Bool, construct_function_without_args::Bool, do_transpose::Bool)
    @ Reactant.TracedUtils /mnt/software/lux/Reactant.jl/src/TracedUtils.jl:197
 [17] make_mlir_fn
    @ /mnt/software/lux/Reactant.jl/src/TracedUtils.jl:117 [inlined]
 [18] #10
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:304 [inlined]
 [19] block!(f::Reactant.Compiler.var"#10#15"{typeof(Enzyme.gradient), Tuple{…}}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201
 [20] #9
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:303 [inlined]
 [21] mmodule!(f::Reactant.Compiler.var"#9#14"{…}, blk::Reactant.MLIR.IR.Module)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Module.jl:92
 [22] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}; optimize::Bool)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:300
 [23] compile_mlir!
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:299 [inlined]
 [24] #6
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:290 [inlined]
 [25] context!(f::Reactant.Compiler.var"#6#7"{@Kwargs{…}, typeof(Enzyme.gradient), Tuple{…}}, ctx::Reactant.MLIR.IR.Context)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Context.jl:76
 [26] compile_mlir(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:288
 [27] top-level scope
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:521
Some type information was truncated. Use `show(err)` to see complete types.

[avik-pal]

julia> st = StateReturn(nothing); @jit simple_ad_fail(x_ra, st)
2-element Vector{Expr}:
 :(result = ConcreteRNumber{Float64}(concrete_res_1))
 :((traced_getfield(args[2], $(Expr(:quote, 1)))).data = concrete_res_2)
:(function var"##simple_ad_fail_reactant#1535336"(args...)
      #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:809 =#
      #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:810 =#
      nothing
      #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:817 =#
      usbuf_1 = (getindex(args, 1)).data
      sbuf_1 = XLA.synced_buffer(usbuf_1)
      #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:818 =#
      begin
          #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:721 =#
          #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:721 =# GC.@preserve sbuf_1 begin
                  #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:722 =#
                  linearized_results = XLA.ExecutableCall(Reactant.XLA.LoadedExecutable(Ptr{Nothing} @0x0000000091fe6490), (sbuf_1.buffer,), (0x01,), Val(2))
              end
          #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:729 =#
          concrete_res_1 = linearized_results[1]
          concrete_res_2 = linearized_results[2]
      end
      #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:819 =#
      ()
      #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:820 =#
      result = ConcreteRNumber{Float64}(concrete_res_1)
      (traced_getfield(args[2], $(Expr(:quote, 1)))).data = concrete_res_2
      #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:821 =#
      return result
  end)
ERROR: type TracedRArray has no field data
Stacktrace:
 [1] setproperty!(x::Reactant.TracedRArray{Float64, 1}, f::Symbol, v::Reactant.XLA.AsyncBuffer)
   @ Base ./Base.jl:51
 [2] macro expansion
   @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:820 [inlined]
 [3] (::Reactant.Compiler.Thunk{Symbol("##simple_ad_fail_reactant#1535336")})(::ConcreteRArray{Float64, 1}, ::StateReturn)
   @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:837
 [4] top-level scope
   @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:475