Fix forward argument indexing

src/Enzyme.jl

@@ -449,6 +449,18 @@ function overload_autodiff(
             push!(
                 outtys, TracedUtils.transpose_ty(MLIR.IR.type(TracedUtils.get_mlir_data(a)))
             )
+
+            if CMode <: ForwardMode && act != enzyme_const && act != enzyme_constnoneed
+                push!(
+                    outtys,
+                    TracedUtils.batch_ty(
+                        width,
+                        TracedUtils.transpose_ty(
+                            MLIR.IR.type(TracedUtils.get_mlir_data(a))
+                        ),
+                    ),
+                )
+            end
         end
     end
 
@@ -517,9 +529,17 @@ function overload_autodiff(
                 arg.val, path[3:end], TracedUtils.transpose_val(MLIR.IR.result(res, residx))
             )
             residx += 1
+            act = act_from_type(arg, reverse, true)
+            if CMode <: ForwardMode && act != enzyme_const && act != enzyme_constnoneed
+                residx += 1
+            end
         else
             TracedUtils.set!(a, (), TracedUtils.transpose_val(MLIR.IR.result(res, residx)))
             residx += 1
+            act = act_from_type(Const, reverse, true)
+            if CMode <: ForwardMode && act != enzyme_const && act != enzyme_constnoneed
+                residx += 1
+            end
         end
     end
 

test/core/autodiff.jl

@@ -646,3 +646,14 @@ end
     model = TwoArgStruct(ConcreteRNumber(0), Reactant.to_rarray(ones(Float32, 3)))
     @jit differentiate_two_arg_fn(model)
 end
+
+function return_arg(x)
+    return x
+end
+
+@testset "Return Argument ForwardMode" begin
+    x = Reactant.to_rarray([1.0, 2.0])
+    dx = Reactant.to_rarray([3.0, 4.0])
+    res = @jit Enzyme.autodiff(Forward, return_arg, Duplicated(x, dx))
+    @test res[1] ≈ dx
+end

test/nn/lux.jl

@@ -97,3 +97,18 @@ end
     res, st_new = @jit model(x, ps, st)
     @test st_new.rng isa Reactant.ReactantRNG
 end
+@testset "Lux Parameter JVP with Bias" begin
+    noisy = Reactant.TestUtils.construct_test_array(Float32, 5, 2)
+    model = Dense(5 => 4; use_bias=true)
+    ps, st = Lux.setup(Xoshiro(0), model)
+    cnoisy = Reactant.to_rarray(noisy)
+    cps = Reactant.to_rarray(ps)
+    cst = Reactant.to_rarray(st)
+
+    sm_input = Lux.StatefulLuxLayer{true}(model, cps, cst)
+    f(p) = first(model(cnoisy, p, cst))
+
+    # test jvp
+    jvp = @jit jacobian_vector_product(f, AutoEnzyme(), cps, cps)
+    @test jvp isa NamedTuple
+end