Preserve residual shape in analytic JVPs

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/SciMLJacobianOperators/Project.toml

@@ -1,6 +1,6 @@
 name = "SciMLJacobianOperators"
 uuid = "19f34311-ddf3-4b8b-af20-060888a46c0e"
-version = "0.1.15"
+version = "0.1.16"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]

lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl

@@ -387,7 +387,7 @@ function prepare_jvp(
                 return
             end
         else
-            return @closure (v, u, p) -> reshape(f.jac(u, p) * vec(v), size(u))
+            return @closure (v, u, p) -> reshape(f.jac(u, p) * vec(v), size(fu))
         end
     end
 

lib/SciMLJacobianOperators/test/core_tests__item3.jl

@@ -104,3 +104,23 @@ prob = NonlinearProblem(
         @test JᵀJv ≈ JᵀJv_analytic atol = 1.0e-5
     end
 end
+
+rectangular_residual(u, p) = reshape(
+    [u[1] + u[2], 2 * u[1] - u[2], u[1] - 3 * u[2]], 3, 1
+)
+rectangular_jacobian(u, p) = [1 1; 2 -1; 1 -3]
+rectangular_u = [2.0, 1.0]
+rectangular_fu = rectangular_residual(rectangular_u, nothing)
+rectangular_prob = NonlinearLeastSquaresProblem(
+    NonlinearFunction{false}(rectangular_residual; jac = rectangular_jacobian), rectangular_u
+)
+
+@testset "Rectangular Analytic Jacobian" begin
+    jac_op = JacobianOperator(rectangular_prob, rectangular_fu, rectangular_u)
+    sop = StatefulJacobianOperator(jac_op, rectangular_u, rectangular_prob.p)
+    v = [3.0, 2.0]
+
+    Jv = sop * v
+    @test size(Jv) == size(rectangular_fu)
+    @test Jv ≈ reshape(rectangular_jacobian(rectangular_u, nothing) * v, 3, 1)
+end