Using scan is a good idea - thanks 
I’ll have to try a few things tomorrow. There’s potentially precedent here but I need to work through it to understand more…
At the moment it’s not especially obvious to me why gradient.jacobian(expression, wrt, ... expects expression as a 1D, rather than handing larger dims. Possibly combinatoric expansion?