test: add #416 as a test

EnzymeAD · Dec 23, 2024 · 5afd2e9 · 5afd2e9
1 parent b969893
commit 5afd2e9
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/src/Compiler.jl b/src/Compiler.jl
@@ -635,7 +635,6 @@ function codegen_unflatten!(
                 unflatcode = :(args[$(path[2])])
                 path = path[3:end]
             else
-                @show "path[1] == $(path[1]) has been ignored..." # XXX: Validate this is correct
                 continue
             end
 

diff --git a/test/autodiff.jl b/test/autodiff.jl
@@ -75,3 +75,25 @@ end
     @test typeof(res) == Tuple{Enzyme.TupleArray{ConcreteRNumber{Float64},(2, 2),4,2}}
     @test res[1] ≈ ones(2, 2)
 end
+
+mutable struct StateReturn
+    st::Any
+end
+
+function cached_return(x, stret::StateReturn)
+    loss = sum(x)
+    stret.st = x .+ 1
+    return loss
+end
+
+@testset "Cached Return: Issue #416" begin
+    x = rand(10)
+    x_ra = Reactant.to_rarray(x)
+
+    stret = StateReturn(nothing)
+    ret = @jit Enzyme.gradient(Reverse, cached_return, x_ra, Const(stret))
+
+    @test @allowscalar all(isone, ret[1])
+    @test stret.st isa ConcreteRArray
+    @test stret.st ≈ x .+ 1
+end