fix: handle constant ConstantOrCache with Enzyme

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.47"
+version = "0.6.48"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -54,15 +54,27 @@ force_annotation(f::F) where {F} = Const(f)
 end
 
 @inline function _translate(
-    backend::AutoEnzyme, ::Mode, ::Val{B}, c::Union{DI.Cache,DI.GeneralizedConstantOrCache}
+    backend::AutoEnzyme, mode::Mode, ::Val{B}, c::DI.Cache
 ) where {B}
+    # important to keep make_zero here for ConstantOrCache instead of similar
     if B == 1
         return Duplicated(DI.unwrap(c), make_zero(DI.unwrap(c)))
     else
         return BatchDuplicated(DI.unwrap(c), ntuple(_ -> make_zero(DI.unwrap(c)), Val(B)))
     end
 end
 
+@inline function _translate(
+    backend::AutoEnzyme, mode::Mode, valB::Val{B}, c::DI.GeneralizedConstantOrCache
+) where {B}
+    IA = guess_activity(typeof(DI.unwrap(c)), mode)
+    if IA <: Const
+        return _translate(backend, mode, valB, DI.Constant(DI.unwrap(c)))
+    else
+        return _translate(backend, mode, valB, DI.Cache(DI.unwrap(c)))
+    end
+end
+
 @inline function _translate(
     backend::AutoEnzyme, mode::Mode, ::Val{B}, c::DI.FunctionContext
 ) where {B}

DifferentiationInterface/test/Back/Enzyme/test.jl

@@ -136,3 +136,14 @@ end
         logging=LOGGING,
     )
 end
+
+@testset "Coverage" begin
+    # ConstantOrCache without cache
+    f_nocontext(x, p) = x
+    @test I == DifferentiationInterface.jacobian(
+        f_nocontext, AutoEnzyme(; mode=Enzyme.Forward), rand(10), ConstantOrCache(nothing)
+    )
+    @test I == DifferentiationInterface.jacobian(
+        f_nocontext, AutoEnzyme(; mode=Enzyme.Reverse), rand(10), ConstantOrCache(nothing)
+    )
+end