Jax transform bug

[bonneted]

The previous update to handle transform with Jax was not working to compute the Hessian. In that case, the output of the Jacobian is reused to compute the Hessian. There was an index error because squeeze make a one element array 0D but we actually want to keep one dimension. Changing the reshape(-1) solved it :

function       |  x.shape
.squeeze:        (1,1) --> ()
.reshape(-1):    (1,1) --> (1,)

I also implemented Jax on Helmotz equation example