Patch bug

LuxDL · Mar 13, 2024 · 978a6c3 · 978a6c3
1 parent 3b2654e
commit 978a6c3
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/ext/BatchedRoutinesFiniteDiffExt.jl b/ext/BatchedRoutinesFiniteDiffExt.jl
@@ -32,7 +32,9 @@ end
 
 # NOTE: This doesn't exploit batching
 @inline function BatchedRoutines._batched_gradient(ad::AutoFiniteDiff, f::F, x) where {F}
-    return FiniteDiff.finite_difference_batched_gradient(f, x, ad.fdjtype)
+    return FiniteDiff.finite_difference_gradient(f, x, ad.fdjtype)
 end
 
+# TODO: For the gradient call just use FiniteDiff over FiniteDiff
+
 end
diff --git a/ext/BatchedRoutinesForwardDiffExt.jl b/ext/BatchedRoutinesForwardDiffExt.jl
@@ -162,8 +162,8 @@ function BatchedRoutines._jacobian_vector_product(ad::AutoForwardDiff, f::F, x,
     Tag = ad.tag === nothing ? typeof(ForwardDiff.Tag(f, eltype(x))) : typeof(ad.tag)
     T = promote_type(eltype(x), eltype(u))
     partials = ForwardDiff.Partials{1, T}.(tuple.(u))
-    u_dual = ForwardDiff.Dual{Tag, T, 1}.(u, partials)
-    y_dual = f(u_dual)
+    x_dual = ForwardDiff.Dual{Tag, T, 1}.(x, partials)
+    y_dual = f(x_dual)
     return ForwardDiff.partials.(y_dual, 1)
 end
 
@@ -172,8 +172,8 @@ function BatchedRoutines._jacobian_vector_product(
     Tag = ad.tag === nothing ? typeof(ForwardDiff.Tag(f, eltype(x))) : typeof(ad.tag)
     T = promote_type(eltype(x), eltype(u))
     partials = ForwardDiff.Partials{1, T}.(tuple.(u))
-    u_dual = ForwardDiff.Dual{Tag, T, 1}.(u, partials)
-    y_dual = f(u_dual, p)
+    x_dual = ForwardDiff.Dual{Tag, T, 1}.(x, partials)
+    y_dual = f(x_dual, p)
     return ForwardDiff.partials.(y_dual, 1)
 end
 

diff --git a/ext/BatchedRoutinesReverseDiffExt.jl b/ext/BatchedRoutinesReverseDiffExt.jl
@@ -18,6 +18,8 @@ function BatchedRoutines._batched_gradient(::AutoReverseDiff, f::F, u) where {F}
     return ReverseDiff.gradient(f, u)
 end
 
+# TODO: Fix the gradient call over ReverseDiff
+
 @concrete struct ReverseDiffPullbackFunction <: Function
     tape
     ∂input