Handle closures with Enzyme

DifferentiationInterface/Project.toml

@@ -44,7 +44,7 @@ DifferentiationInterfaceTrackerExt = "Tracker"
 DifferentiationInterfaceZygoteExt = ["Zygote", "ForwardDiff"]
 
 [compat]
-ADTypes = "1.5.0"
+ADTypes = "1.6.1"
 ChainRulesCore = "1.23.0"
 Compat = "3,4"
 Diffractor = "=0.2.6"

DifferentiationInterface/docs/src/backends.md

@@ -57,7 +57,7 @@ import Zygote
 
 backend_examples = [
     AutoDiffractor(),
-    AutoEnzyme(),
+    AutoEnzyme(; constant_function=true),
     AutoFastDifferentiation(),
     AutoFiniteDiff(),
     AutoFiniteDifferences(; fdm=FiniteDifferences.central_fdm(3, 1)),

DifferentiationInterface/docs/src/tutorial1.md

@@ -116,7 +116,7 @@ Typically, for gradients, reverse mode AD might be a better fit, so let's try th
 ```@example tuto1
 import Enzyme
 
-backend2 = AutoEnzyme()
+backend2 = AutoEnzyme(constant_function=true)
 ```
 
 Once the backend is created, things run smoothly with exactly the same syntax as before:

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -38,15 +38,19 @@ using Enzyme:
     make_zero,
     make_zero!
 
-struct AutoDeferredEnzyme{M} <: ADTypes.AbstractADType
+struct AutoDeferredEnzyme{M,constant_function} <: ADTypes.AbstractADType
     mode::M
 end
 
 ADTypes.mode(backend::AutoDeferredEnzyme) = ADTypes.mode(AutoEnzyme(backend.mode))
 
-DI.nested(backend::AutoEnzyme) = AutoDeferredEnzyme(backend.mode)
+function DI.nested(backend::AutoEnzyme{M,constant_function}) where {M,constant_function}
+    return AutoDeferredEnzyme{M,constant_function}(backend.mode)
+end
 
-const AnyAutoEnzyme{M} = Union{AutoEnzyme{M},AutoDeferredEnzyme{M}}
+const AnyAutoEnzyme{M,constant_function} = Union{
+    AutoEnzyme{M,constant_function},AutoDeferredEnzyme{M,constant_function}
+}
 
 # forward mode if possible
 forward_mode(backend::AnyAutoEnzyme{<:Mode}) = backend.mode
@@ -68,6 +72,15 @@ function DI.basis(::AutoEnzyme, a::AbstractArray{T}, i::CartesianIndex) where {T
     return b
 end
 
+function get_f_and_df(f, ::AnyAutoEnzyme{M,true}) where {M}
+    return Const(f)
+end
+
+function get_f_and_df(f, ::AnyAutoEnzyme{M,false}) where {M}
+    df = make_zero(f)
+    return Duplicated(f, df)
+end
+
 include("forward_onearg.jl")
 include("forward_twoarg.jl")
 

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -7,25 +7,27 @@ end
 function DI.value_and_pushforward(
     f, backend::AnyAutoEnzyme{<:Union{ForwardMode,Nothing}}, x, dx, ::NoPushforwardExtras
 )
+    f_and_df = get_f_and_df(f, backend)
     dx_sametype = convert(typeof(x), dx)
     x_and_dx = Duplicated(x, dx_sametype)
     y, new_dy = if backend isa AutoDeferredEnzyme
-        autodiff_deferred(forward_mode(backend), f, Duplicated, x_and_dx)
+        autodiff_deferred(forward_mode(backend), f_and_df, Duplicated, x_and_dx)
     else
-        autodiff(forward_mode(backend), Const(f), Duplicated, x_and_dx)
+        autodiff(forward_mode(backend), f_and_df, Duplicated, x_and_dx)
     end
     return y, new_dy
 end
 
 function DI.pushforward(
     f, backend::AnyAutoEnzyme{<:Union{ForwardMode,Nothing}}, x, dx, ::NoPushforwardExtras
 )
+    f_and_df = get_f_and_df(f, backend)
     dx_sametype = convert(typeof(x), dx)
     x_and_dx = Duplicated(x, dx_sametype)
     new_dy = if backend isa AutoDeferredEnzyme
-        only(autodiff_deferred(forward_mode(backend), f, DuplicatedNoNeed, x_and_dx))
+        only(autodiff_deferred(forward_mode(backend), f_and_df, DuplicatedNoNeed, x_and_dx))
     else
-        only(autodiff(forward_mode(backend), Const(f), DuplicatedNoNeed, x_and_dx))
+        only(autodiff(forward_mode(backend), f_and_df, DuplicatedNoNeed, x_and_dx))
     end
     return new_dy
 end
@@ -61,34 +63,42 @@ struct EnzymeForwardGradientExtras{B,O} <: GradientExtras
     shadow::O
 end
 
-function DI.prepare_gradient(f, backend::AutoEnzyme{<:ForwardMode}, x)
+function DI.prepare_gradient(f, backend::AutoEnzyme{<:ForwardMode,true}, x)
     B = pick_batchsize(backend, length(x))
     shadow = chunkedonehot(x, Val(B))
     return EnzymeForwardGradientExtras{B,typeof(shadow)}(shadow)
 end
 
 function DI.gradient(
-    f, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{B}
+    f, backend::AutoEnzyme{<:ForwardMode,true}, x, extras::EnzymeForwardGradientExtras{B}
 ) where {B}
     grad_tup = gradient(forward_mode(backend), f, x, Val(B); shadow=extras.shadow)
     return reshape(collect(grad_tup), size(x))
 end
 
 function DI.value_and_gradient(
-    f, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras
+    f, backend::AutoEnzyme{<:ForwardMode,true}, x, extras::EnzymeForwardGradientExtras
 )
     return f(x), DI.gradient(f, backend, x, extras)
 end
 
 function DI.gradient!(
-    f, grad, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{B}
+    f,
+    grad,
+    backend::AutoEnzyme{<:ForwardMode,true},
+    x,
+    extras::EnzymeForwardGradientExtras{B},
 ) where {B}
     grad_tup = gradient(forward_mode(backend), f, x, Val(B); shadow=extras.shadow)
     return copyto!(grad, grad_tup)
 end
 
 function DI.value_and_gradient!(
-    f, grad, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{B}
+    f,
+    grad,
+    backend::AutoEnzyme{<:ForwardMode,true},
+    x,
+    extras::EnzymeForwardGradientExtras{B},
 ) where {B}
     grad_tup = gradient(forward_mode(backend), f, x, Val(B); shadow=extras.shadow)
     return f(x), copyto!(grad, grad_tup)
@@ -100,15 +110,15 @@ struct EnzymeForwardOneArgJacobianExtras{B,O} <: JacobianExtras
     shadow::O
 end
 
-function DI.prepare_jacobian(f, backend::AutoEnzyme{<:Union{ForwardMode,Nothing}}, x)
+function DI.prepare_jacobian(f, backend::AutoEnzyme{<:Union{ForwardMode,Nothing},true}, x)
     B = pick_batchsize(backend, length(x))
     shadow = chunkedonehot(x, Val(B))
     return EnzymeForwardOneArgJacobianExtras{B,typeof(shadow)}(shadow)
 end
 
 function DI.jacobian(
     f,
-    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing},true},
     x,
     extras::EnzymeForwardOneArgJacobianExtras{B},
 ) where {B}
@@ -120,7 +130,7 @@ end
 
 function DI.value_and_jacobian(
     f,
-    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing},true},
     x,
     extras::EnzymeForwardOneArgJacobianExtras,
 )
@@ -130,7 +140,7 @@ end
 function DI.jacobian!(
     f,
     jac,
-    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing},true},
     x,
     extras::EnzymeForwardOneArgJacobianExtras,
 )
@@ -140,7 +150,7 @@ end
 function DI.value_and_jacobian!(
     f,
     jac,
-    backend::AnyAutoEnzyme{<:Union{ForwardMode,Nothing}},
+    backend::AnyAutoEnzyme{<:Union{ForwardMode,Nothing},true},
     x,
     extras::EnzymeForwardOneArgJacobianExtras,
 )

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -12,14 +12,15 @@ function DI.value_and_pushforward(
     dx,
     ::NoPushforwardExtras,
 )
+    f!_and_df! = get_f_and_df(f!, backend)
     dx_sametype = convert(typeof(x), dx)
     dy_sametype = make_zero(y)
     y_and_dy = Duplicated(y, dy_sametype)
     x_and_dx = Duplicated(x, dx_sametype)
     if backend isa AutoDeferredEnzyme
-        autodiff_deferred(forward_mode(backend), f!, Const, y_and_dy, x_and_dx)
+        autodiff_deferred(forward_mode(backend), f!_and_df!, Const, y_and_dy, x_and_dx)
     else
-        autodiff(forward_mode(backend), Const(f!), Const, y_and_dy, x_and_dx)
+        autodiff(forward_mode(backend), f!_and_df!, Const, y_and_dy, x_and_dx)
     end
     return y, dy_sametype
 end