feat: lower to a looped batched jacobian

Author: avik-pal

ext/LuxReactantExt/batched_jacobian.jl

@@ -22,33 +22,43 @@ function (f::ApplyWithReshape)(y, x)
     return nothing
 end
 
-function _batched_jacobian_reverse_impl(f::F, ad::AutoEnzyme, x::AbstractArray) where {F}
+function _check_validity_for_batched_jacobian(f::F, x::AbstractArray) where {F}
     y = f(x)
     @assert y isa AbstractArray
-    if ndims(y) ≤ 1 || size(y, ndims(y)) != size(x, ndims(x))
+    B = size(y, ndims(y))
+    if ndims(y) ≤ 1 || B != size(x, ndims(x))
         throw(AssertionError("`batched_jacobian` only supports batched outputs \
                               (ndims(y) > 1) && size(y, ndims(y)) == size(x, ndims(x))."))
     end
+    return y, B
+end
 
+function _batched_jacobian_reverse_impl(f::F, ad::AutoEnzyme, x::AbstractArray) where {F}
+    y, B = _check_validity_for_batched_jacobian(f, x)
     f′ = ApplyWithReshape(f, size(x))
 
-    y = Utils.contiguous(reshape(y, :, size(y, ndims(y))))
-    dy = repeat(
-        reshape(
-            Reactant.promote_to(
-                TracedRArray{Reactant.unwrapped_eltype(y),2}, LinearAlgebra.I(size(y, 1))
+    y = Utils.contiguous(reshape(y, :, B))
+    dy = Utils.contiguous(
+        repeat(
+            reshape(
+                Reactant.promote_to(
+                    TracedRArray{Reactant.unwrapped_eltype(y),2},
+                    LinearAlgebra.I(size(y, 1)),
+                ),
+                size(y, 1),
+                1,
+                size(y, 1),
             ),
-            size(y, 1),
             1,
-            size(y, 1),
+            size(y, 2),
+            1,
         ),
-        1,
-        size(y, 2),
-        1,
     )
-    dy = Utils.contiguous(dy)
 
-    x = Utils.contiguous(reshape(x, :, size(x, ndims(x))))
+    x = Utils.contiguous(reshape(x, :, B))
+
+    # TODO: replace once https://github.com/LuxDL/Lux.jl/issues/1523 is fixed
+    #=
     dx = similar(x, size(x, 1), size(x, 2), size(y, 1))
     fill!(dx, false)
 
@@ -60,35 +70,71 @@ function _batched_jacobian_reverse_impl(f::F, ad::AutoEnzyme, x::AbstractArray)
     )
 
     return permutedims(dx, (3, 1, 2))
+    =#
+
+    # Our loop to batch pass should automatically batch this loop and current has better
+    # coverage than the above. Though we should fix the above to ensure we never have a
+    # loop in the final result.
+    dx = similar(x, size(y, 1), size(x, 1), size(x, 2))
+    @trace track_numbers = false for i in 1:size(y, 1)
+        dxᵢ = Enzyme.make_zero(x)
+        Enzyme.autodiff(
+            ad.mode,
+            Utils.annotate_enzyme_function(ad, f′),
+            Duplicated(y, dy[:, :, i]),
+            Duplicated(x, dxᵢ),
+        )
+        dx[i, :, :] = dxᵢ
+    end
+    return dx
 end
 
 function _batched_jacobian_forward_impl(f::F, ad::AutoEnzyme, x::AbstractArray) where {F}
+    y, B = _check_validity_for_batched_jacobian(f, x)
+    y = Utils.contiguous(reshape(y, :, B)) # will be DCEd away
+
     f′ = ApplyWithReshape(f, size(x))
     x = Utils.contiguous(reshape(x, :, size(x, ndims(x))))
 
-    bx = repeat(
-        reshape(
-            Reactant.promote_to(
-                TracedRArray{Reactant.unwrapped_eltype(x),2}, LinearAlgebra.I(size(x, 1))
+    bx = Utils.contiguous(
+        repeat(
+            reshape(
+                Reactant.promote_to(
+                    TracedRArray{Reactant.unwrapped_eltype(x),2},
+                    LinearAlgebra.I(size(x, 1)),
+                ),
+                size(x, 1),
+                1,
+                size(x, 1),
             ),
-            size(x, 1),
             1,
-            size(x, 1),
+            size(x, 2),
+            1,
         ),
-        1,
-        size(x, 2),
-        1,
-    )
-    bx = Utils.contiguous(bx)
-
-    return stack(
-        only(
-            Enzyme.autodiff(
-                ad.mode,
-                Utils.annotate_enzyme_function(ad, f′),
-                Reactant.StackedBatchDuplicated(x, bx),
-            ),
-        );
-        dims=2,
     )
+
+    # TODO: replace once https://github.com/LuxDL/Lux.jl/issues/1523 is fixed
+    # return stack(
+    #     only(
+    #         Enzyme.autodiff(
+    #             ad.mode,
+    #             Utils.annotate_enzyme_function(ad, f′),
+    #             Reactant.StackedBatchDuplicated(x, bx),
+    #         ),
+    #     );
+    #     dims=2,
+    # )
+
+    dy = similar(y, size(y, 1), size(x, 1), size(x, 2))
+    @trace track_numbers = false for i in 1:size(x, 1)
+        dyᵢ = Enzyme.make_zero(y)
+        Enzyme.autodiff(
+            ad.mode,
+            Utils.annotate_enzyme_function(ad, f′),
+            Duplicated(y, dyᵢ),
+            Duplicated(x, bx[:, :, i]),
+        )
+        dy[:, i, :] = dyᵢ
+    end
+    return dy
 end

test/reactant/autodiff_tests.jl

@@ -59,7 +59,7 @@ end
 @testitem "AutoDiff APIs: Batched Jacobian" tags = [:reactant] setup = [SharedTestSetup] begin
     using Reactant, Lux, Zygote, Random, ForwardDiff, Enzyme
 
-    fn(x) = reshape(sum(abs2, x; dims=(1, 2, 3)), 1, :)
+    fn(x) = reshape(sum(abs2, x; dims=(2, 3)), size(x, 1), :)
 
     rng = Random.default_rng()