DuplicatedNoNeed will provide incorrect results if values are read-write like for a cache (#14)

Author: vchuravy

src/NewtonKrylov.jl

@@ -13,7 +13,7 @@ import LinearAlgebra: mul!
 
 function maybe_duplicated(f, df)
     if !Enzyme.Compiler.guaranteed_const(typeof(f))
-        return DuplicatedNoNeed(f, df)
+        return Duplicated(f, df)
     else
         return Const(f)
     end
@@ -37,15 +37,16 @@ end
 
 Base.size(J::JacobianOperator) = (length(J.res), length(J.u))
 Base.eltype(J::JacobianOperator) = eltype(J.u)
+Base.length(J::JacobianOperator) = prod(size(J))
 
 function mul!(out, J::JacobianOperator, v)
     # Enzyme.make_zero!(J.f_cache)
     f_cache = Enzyme.make_zero(J.f) # Stop gap until we can zero out mutable values
     autodiff(
         Forward,
         maybe_duplicated(J.f, f_cache), Const,
-        DuplicatedNoNeed(J.res, reshape(out, size(J.res))),
-        DuplicatedNoNeed(J.u, reshape(v, size(J.u)))
+        Duplicated(J.res, reshape(out, size(J.res))),
+        Duplicated(J.u, reshape(v, size(J.u)))
     )
     return nothing
 end
@@ -61,16 +62,18 @@ function mul!(out, J′::Union{Adjoint{<:Any, <:JacobianOperator}, Transpose{<:A
     Enzyme.make_zero!(J.f_cache)
     # TODO: provide cache for `copy(v)`
     # Enzyme zeros input derivatives and that confuses the solvers.
+    # If `out` is non-zero we might get spurious gradients
+    fill!(out, 0)
     autodiff(
         Reverse,
         maybe_duplicated(J.f, J.f_cache), Const,
-        DuplicatedNoNeed(J.res, reshape(copy(v), size(J.res))),
-        DuplicatedNoNeed(J.u, reshape(out, size(J.u)))
+        Duplicated(J.res, reshape(copy(v), size(J.res))),
+        Duplicated(J.u, reshape(out, size(J.u)))
     )
     return nothing
 end
 
-function Base.collect(JOp::JacobianOperator)
+function Base.collect(JOp::Union{Adjoint{<:Any, <:JacobianOperator}, Transpose{<:Any, <:JacobianOperator}, JacobianOperator})
     N, M = size(JOp)
     v = zeros(eltype(JOp), M)
     out = zeros(eltype(JOp), N)

test/runtests.jl

@@ -28,6 +28,19 @@ using Enzyme, LinearAlgebra
 @testset "Jacobian" begin
     J_Enz = jacobian(Forward, F, [3.0, 5.0]) |> only
     J = JacobianOperator(F!, zeros(2), [3.0, 5.0])
+
+    @test size(J) == (2, 2)
+    @test length(J) == 4
+    @test eltype(J) == Float64
+
+    out = [NaN, NaN]
+    mul!(out, J, [1.0, 0.0])
+    @test out == [6.0, 7.38905609893065]
+
+    out = [NaN, NaN]
+    mul!(out, transpose(J), [1.0, 0.0])
+    @test out == [6.0, 10.0]
+
     J_NK = collect(J)
 
     @test J_NK == J_Enz
@@ -37,4 +50,6 @@ using Enzyme, LinearAlgebra
     mul!(out, J, v)
 
     @test out ≈ J_Enz * v
+
+    @test collect(transpose(J)) == transpose(collect(J))
 end