Update jax.py

Author: heptaflar

uncertaintyx/fit/eiv/jax.py

@@ -57,7 +57,7 @@ def evm(
     max_i: int = DEFAULT_MAX_I,
     max_g: Any = DEFAULT_MAX_G,
     covar: bool = False,
-) -> tuple[Any, ...]:
+) -> tuple[Array, Array, Array, Array, bool]:
     r"""
     Implementation of the effective variance method (EVM) with
     a limited-memory Broyden-Fletcher-Goldfarb-Shanno (L-BFGS)
@@ -255,7 +255,7 @@ def post(p: Array) -> tuple[Array, Array]:
     popt, cost, converged = opti(p)
     pcov, punc = post(popt)
 
-    return popt, pcov, punc, cost, converged
+    return popt, pcov, punc, cost, converged.any()
 
 
 class EIV(Fitting):
@@ -315,5 +315,5 @@ def fit(
             pcov=np.asarray(pcov),
             rvar=rvar,
             cost=np.asarray(cost),
-            info=0 if converged.item() else 1,
+            info=0 if converged else 1,
         )