Fix multiarg fwd and fwd mode copy on 1.11 (#2469)

* Fix multiarg fwd and fwd mode copy on 1.11

* cleanup

Author: wsmoses

src/rules/llvmrules.jl

@@ -805,7 +805,7 @@ end
 
             LLVM.memset!(
                 B,
-                inttoptr!(B, ev2, LLVM.PointerType(LLVM.IntType(8))),
+                get_memory_data(B, callv),
                 LLVM.ConstantInt(i8, 0, false),
                 length,
                 algn,

src/sugar.jl

@@ -487,15 +487,6 @@ end
 	res
 end
 
-@inline function tupstack(x, outshape::Tuple{Vararg{Int}}, inshape::Tuple{Vararg{Int}})
-    st = Base.stack(x)
-    if length(outshape) == 1
-        st
-    else
-        reshape(st, (outshape..., inshape...))
-    end
-end
-
 @inline specialize_output(output, input) = output
 
 """
@@ -744,12 +735,14 @@ gradient(Forward, mul, [2.0, 3.0], Const([2.7, 3.1]))
                 if argnum > 0
                     quote
                         if $tmp[1] isa AbstractArray
-                            inshape = size($(vals[1]))
+                            inshape = size($(vals[i]))
                             outshape = size($tmp[1])
+                            num = prod(outshape)
+
                             # st : outshape x total inputs
                             tupstack($tmp, outshape, inshape)
                         else
-                            specialize_output(TupleArray($tmp, size($arg)), $(vals[1]))
+                            specialize_output(TupleArray($tmp, size($arg)), $(vals[i]))
                         end
                     end
                 else

test/sugar.jl

@@ -10,6 +10,12 @@ mul_vector(x, y) = [x[1]*y[2], x[2]*y[1]]
     @test res[1] === inp
 end
 
+
+function diffsize(θ0, X)
+    return copy(X)
+end
+
+
 @testset "Forward Multi-Arg Gradient" begin
 	res = gradient(Forward, mul_scalar, [2.0, 3.0], [2.7, 3.1])
 	@test res[1] ≈ [3.1, 2.7]
@@ -38,7 +44,9 @@ end
 	@test res.derivs[1] ≈ [3.1, 2.7]
 	@test res.derivs[2] ≈ [3.0, 2.0]
 
-
+    res = gradient(Forward, diffsize, [2.0, 3.0], [2.7, 3.1, 4.5])
+    @test res[1] ≈ zeros(3, 2)
+    @test res[2] ≈ [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]
 
 	res = gradient(Forward, mul_scalar, Const([2.0, 3.0]), [2.7, 3.1])
 	@test res[1] == nothing