Wrong JVP and VJP ?

[NightWinkle]

Hi,

I am trying to compute VJPs and JVPs for a simple heat equation, but it seems unless I am missing something, the output shapes are inconsistent :

function heat(u, p, t)
    gamma, dx = p
    dxxu = (circshift(u, 1) .- 2 * u .+ circshift(u, -1)) / dx^2
    return gamma * dxxu
end

Ns = 5
xmin, xmax = 0.0, 1.0
dx = (xmax - xmin) / (Ns - 1)

xs = collect(LinRange(xmin, xmax, Ns))
using FastDifferentiation
FastDifferentiation.@variables uv pv tv
uvs = make_variables(:uv, Ns)
pvs = make_variables(:pv, 2)
tvs = make_variables(:tv, 1)
fv = heat(uvs, pvs, tvs)
js = FastDifferentiation.jacobian(fv, uvs) # The Jacobian js is of shape Ns x Ns

vjps, rvs = FastDifferentiation.jacobian_transpose_v(fv, uvs) # The variable vjps is of shape Ns+2 and rvs is of shape Ns
jvps, vvs = FastDifferentiation.jacobian_times_v(fv, uvs) # The variable jvps is of shape Ns and rvs is of shape Ns+2

The jacobian seems correct but not the vjps and jvps.

[brianguenter]

Yes, it's a bug. The vjps terms has space for all the variables in the function, not just the variables you are differentiating with respect to, in this case uvs. Same thing for the vvs, term it has space for all the variables in the functions. I'll fix it today.

[brianguenter]

I have a fix on branch brianguenter/issue60. Could you download this branch and run it on your code to verify that it is giving the results you expect? If you have any other code that exercises jacobian_times_v and jacobian_transpose_v could you run that as well and report any bugs to me?

[brianguenter]

Since I haven't heard from you will assume you are no longer having problems. I'm closing the issue.

[NightWinkle]

Sorry for the late response.

Thanks a lot for your work and for this library ! It seems to be working perfectly.