Update jax.py

Author: heptaflar

uncertaintyx/fit/eiv/jax.py

@@ -243,7 +243,7 @@ def post(p: Array) -> tuple[Array, Array]:
         """
         hess = jax.hessian(S)
         pcov = jli.pinv(hess(p).reshape(p.size, -1))
-        punc = jnp.sqrt(jnp.diag(pcov)).reshape(p.shape)
+        punc = jnp.sqrt(jnp.diag(pcov))
         return pcov.reshape(p.shape + p.shape), punc.reshape(p.shape)
 
     if ux is None: