Improve Enzyme batch size and gradient (#557)

* Adaptive Enzyme batch size

* Use Enzyme.gradient whenever possible

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.8"
+version = "0.6.9"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -219,14 +219,8 @@ function DI.gradient(
     contexts::Vararg{Context,C},
 ) where {F,C}
     f_and_df = get_f_and_df(f, backend)
-    grad = make_zero(x)
-    autodiff(
-        reverse_noprimal(backend),
-        f_and_df,
-        Active,
-        Duplicated(x, grad),
-        map(translate, contexts)...,
-    )
+    ders = gradient(reverse_noprimal(backend), f_and_df, x, map(translate, contexts)...)
+    grad = first(ders)
     return grad
 end
 
@@ -237,14 +231,10 @@ function DI.value_and_gradient(
     contexts::Vararg{Context,C},
 ) where {F,C}
     f_and_df = get_f_and_df(f, backend)
-    grad = make_zero(x)
-    _, y = autodiff(
-        reverse_withprimal(backend),
-        f_and_df,
-        Active,
-        Duplicated(x, grad),
-        map(translate, contexts)...,
+    ders, y = gradient(
+        reverse_withprimal(backend), f_and_df, x, map(translate, contexts)...
     )
+    grad = first(ders)
     return y, grad
 end
 
@@ -272,13 +262,8 @@ function DI.gradient(
     contexts::Vararg{Context,C},
 ) where {F,C}
     f_and_df = get_f_and_df(f, backend)
-    grad = make_zero(x)
-    autodiff(
-        reverse_noprimal(backend),
-        f_and_df,
-        Duplicated(x, grad),
-        map(translate, contexts)...,
-    )
+    ders = gradient(reverse_noprimal(backend), f_and_df, x, map(translate, contexts)...)
+    grad = first(ders)
     return grad
 end
 
@@ -300,7 +285,7 @@ function DI.gradient!(
         Duplicated(x, grad_righttype),
         map(translate, contexts)...,
     )
-    grad isa typeof(x) || copyto!(grad, grad_righttype)
+    grad === grad_righttype || copyto!(grad, grad_righttype)
     return grad
 end
 
@@ -312,14 +297,10 @@ function DI.value_and_gradient(
     contexts::Vararg{Context,C},
 ) where {F,C}
     f_and_df = get_f_and_df(f, backend)
-    grad = make_zero(x)
-    _, y = autodiff(
-        reverse_withprimal(backend),
-        f_and_df,
-        Active,
-        Duplicated(x, grad),
-        map(translate, contexts)...,
+    ders, y = gradient(
+        reverse_withprimal(backend), f_and_df, x, map(translate, contexts)...
     )
+    grad = first(ders)
     return y, grad
 end
 
@@ -341,7 +322,7 @@ function DI.value_and_gradient!(
         Duplicated(x, grad_righttype),
         map(translate, contexts)...,
     )
-    grad isa typeof(x) || copyto!(grad, grad_righttype)
+    grad === grad_righttype || copyto!(grad, grad_righttype)
     return y, grad
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -1,5 +1,5 @@
 # until https://github.com/EnzymeAD/Enzyme.jl/pull/1545 is merged
-DI.pick_batchsize(::AutoEnzyme, dimension::Integer) = Val(16)
+DI.pick_batchsize(::AutoEnzyme, dimension::Integer) = Val(min(dimension, 16))
 
 ## Annotations