Nonsquare-jacobian MultiField autodiff

I used to be able to compute the derivative of the residual with respect to parameters with a small twick of the jacobian function within Gridap the old solution under Gridap v0.15. I thought I could upgrade to Gridap 0.17 to catch up on your amazing progresses and give some help with the autodiff part.

Unfortunately, I don't see any way through and would appreciate some pointers if you have the time. I believe this would greatly benefit and extend the capacities of Gridap in a new direction.

I tried to come up with a simple example:

using Gridap

model = CartesianDiscreteModel((0, 1), 10)
labels = get_face_labeling(model)
add_tag_from_tags!(labels, "diri0", [1]) #initial time condition

valuetype = Float64
reffe = ReferenceFE(lagrangian, valuetype, 1)

Vcon = TestFESpace(model, reffe, conformity = :L2)
Ucon = TrialFESpace(Vcon)
VI = TestFESpace(model, reffe; conformity = :H1, labels = labels, dirichlet_tags = ["diri0"])
UI = TrialFESpace(VI, 0.6)
VS = TestFESpace(model, reffe; conformity = :H1, labels = labels, dirichlet_tags = ["diri0"])
US = TrialFESpace(VS, 0.1)

Xpde = MultiFieldFESpace([VI, VS])
Ypde = MultiFieldFESpace([UI, US])
Xcon = MultiFieldFESpace([Vcon])
Ycon = MultiFieldFESpace([Ucon])

trian = Triangulation(model)
degree = 1
dΩ = Measure(trian, degree)

kp(x) = 1.01
kr(x) = 2.03
function res(y, u, v)
  cf, pf = y
  p, q = v
  ∫(
    p * (kp * pf * (1.0 - cf)) + q * (u[1] * kr * cf * (1.0 - cf - pf))
  )dΩ
end
nv = Gridap.FESpaces.num_free_dofs(Ypde)
yh = FEFunction(Ypde, rand(nv))

nc = Gridap.FESpaces.num_free_dofs(Ycon)
uh = FEFunction(Ycon, rand(nc))

v = Gridap.FESpaces.get_fe_basis(Xpde)
resuh = res(yh, uh, v)

jac_y = Gridap.FESpaces._jacobian(x -> res(x, uh, v), yh, resuh) # works !!

jac_u = Gridap.FESpaces._jacobian(x -> res(yh, x, v), uh, resuh) # fails because 'A check failed'

The error comes from _change_argument(jacobian,f,trian,uh) https://github.com/gridap/Gridap.jl/blob/1fb3dc9abf8c47685637901bd14a74e4355a9492/src/MultiField/MultiFieldFEAutodiff.jl#L73 and https://github.com/gridap/Gridap.jl/blob/1fb3dc9abf8c47685637901bd14a74e4355a9492/src/Arrays/Autodiff.jl#L13

gridap / Gridap.jl

Nonsquare-jacobian MultiField autodiff #774