Replace uses of internal JAX NumPy utils with public API functions. (#…

…487)

* Replace uses of internal JAX NumPy utils with public API functions.

* Bug fix to update to new jaxopt version

blackjax/optimizers/lbfgs.py

@@ -120,6 +120,7 @@ def minimize_lbfgs(
         state=LbfgsState(
             iter_num=last_step_raveled.state.iter_num,
             value=last_step_raveled.state.value,
+            grad=unravel_fn(last_step_raveled.state.grad),
             stepsize=last_step_raveled.state.stepsize,
             error=last_step_raveled.state.error,
             s_history=unravel_fn_mapped(last_step_raveled.state.s_history),

blackjax/util.py

@@ -4,7 +4,6 @@
 
 import jax.numpy as jnp
 from jax import jit, lax
-from jax._src.numpy.util import _promote_dtypes
 from jax.flatten_util import ravel_pytree
 from jax.random import normal
 from jax.tree_util import tree_leaves
@@ -42,7 +41,9 @@ def linear_map(diag_or_dense_a, b, *, precision="highest"):
     -------
         The result vector of the matrix multiplication.
     """
-    diag_or_dense_a, b = _promote_dtypes(diag_or_dense_a, b)
+    dtype = jnp.result_type(diag_or_dense_a.dtype, b.dtype)
+    diag_or_dense_a = diag_or_dense_a.astype(dtype)
+    b = b.astype(dtype)
     ndim = jnp.ndim(diag_or_dense_a)
 
     if ndim <= 1: