Remove deprecated JAX funtions (#5203)

* Remove deprecated JAX funtions

* nox changes

* Try to temporarily use pybammsolvers from git branch

* Try to make temporary fix for CI work

* Revert non-JAX stuff, remove version pin

* Try to make work again with old Jax versions on MacOS-13

* style: pre-commit fixes

* Try with Jax 0.6 for now

* Add unit test for split_list

* style: pre-commit fixes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: DavidMStraub

pyproject.toml

@@ -113,7 +113,7 @@ dev = [
     "hypothesis",
 ]
 jax = [
-    "jax>=0.4.36,<0.6.0",
+    "jax>=0.4.36,<0.7.0",
 ]
 # Contains all optional dependencies, except for jax, and dev dependencies
 all = [

src/pybamm/solvers/idaklu_jax.py

@@ -16,12 +16,22 @@
 if pybamm.has_jax():
     import jax
     from jax import lax
+
+    try:
+        from jax import ffi
+    except ImportError:
+        from jax.extend import ffi
     from jax import numpy as jnp
     from jax.interpreters import ad, batching, mlir
     from jax.interpreters.mlir import custom_call
-    from jax.lib import xla_client
     from jax.tree_util import tree_flatten
 
+    # Handle JAX version compatibility for Primitive location
+    try:
+        from jax.core import Primitive
+    except ImportError:
+        from jax.extend.core import Primitive
+
 
 class IDAKLUJax:
     """JAX wrapper for IDAKLU solver
@@ -600,15 +610,14 @@ def _jaxify(
         self._register_callbacks()  # Register python methods as callbacks in IDAKLU-JAX
 
         for _name, _value in idaklu.registrations().items():
-            # todo: This has been removed from jax v0.6.0
-            xla_client.register_custom_call_target(
-                f"{_name}_{self._unique_name()}", _value, platform="cpu"
+            ffi.register_ffi_target(
+                f"{_name}_{self._unique_name()}", _value, platform="cpu", api_version=0
             )
 
         # --- JAX PRIMITIVE DEFINITION ------------------------------------------------
 
         logger.debug(f"Creating new primitive: {self._unique_name()}")
-        f_p = jax.core.Primitive(f"f_{self._unique_name()}")
+        f_p = Primitive(f"f_{self._unique_name()}")
         f_p.multiple_results = False  # Returns a single multi-dimensional array
 
         def f(t, inputs=None):
@@ -759,7 +768,7 @@ def make_zero(prim, tan):
 
         ad.primitive_jvps[f_p] = f_jvp
 
-        f_jvp_p = jax.core.Primitive(f"f_jvp_{self._unique_name()}")
+        f_jvp_p = Primitive(f"f_jvp_{self._unique_name()}")
 
         @f_jvp_p.def_impl
         def f_jvp_eval(*args):
@@ -941,7 +950,7 @@ def f_jvp_lowering_cpu(ctx, *args):
 
         # --- JAX PRIMITIVE VJP DEFINITION --------------------------------------------
 
-        f_vjp_p = jax.core.Primitive(f"f_vjp_{self._unique_name()}")
+        f_vjp_p = Primitive(f"f_vjp_{self._unique_name()}")
 
         def f_vjp(y_bar, invar, *primals):
             """Main wrapper for the VJP function"""

src/pybamm/solvers/jax_bdf_solver.py

@@ -8,6 +8,8 @@
 import pybamm
 
 if pybamm.has_jax():
+    import functools
+
     import jax
     import jax.numpy as jnp
     from jax import core, dtypes
@@ -16,7 +18,19 @@
     from jax.flatten_util import ravel_pytree
     from jax.interpreters import partial_eval as pe
     from jax.tree_util import tree_flatten, tree_map, tree_unflatten
-    from jax.util import cache, safe_map, split_list
+
+    def split_list(lst, indices):
+        """Split a list at given indices."""
+        if not indices:
+            return [lst]
+
+        result = []
+        start = 0
+        for idx in indices:
+            result.append(lst[start:idx])
+            start = idx
+        result.append(lst[start:])
+        return result
 
     platform = jax.lib.xla_bridge.get_backend().platform.casefold()
     if platform != "metal":
@@ -875,7 +889,7 @@ def scan_fun(carry, i):
 
     _bdf_odeint.defvjp(_bdf_odeint_fwd, _bdf_odeint_rev)
 
-    @cache()
+    @functools.cache
     def closure_convert(fun, in_tree, in_avals):
         wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
         jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
@@ -975,7 +989,7 @@ def _check_arg(arg):
         raise TypeError(msg.format(arg))
 
     flat_args, in_tree = tree_flatten((y0, t_eval[0], *args))
-    in_avals = tuple(safe_map(abstractify, flat_args))
+    in_avals = tuple(map(abstractify, flat_args))
     converted, consts = closure_convert(func, in_tree, in_avals)
     if mass is None:
         mass = onp.identity(y0.shape[0], dtype=y0.dtype)

tests/unit/test_solvers/test_jax_bdf_solver.py

@@ -175,3 +175,36 @@ def fun(y, t, inputs):
         )
 
         np.testing.assert_allclose(y[:, 0].reshape(-1), np.exp(-0.1 * t_eval))
+
+    def test_split_list(self):
+        """Test the split_list utility function."""
+        from pybamm.solvers.jax_bdf_solver import split_list
+
+        # Test case 1: Empty indices should return the original list
+        original_list = [1, 2, 3, 4, 5]
+        result = split_list(original_list, [])
+        assert result == [[1, 2, 3, 4, 5]]
+
+        # Test case 2: Single index should split at that point
+        result = split_list([1, 2, 3, 4, 5], [3])
+        assert result == [[1, 2, 3], [4, 5]]
+
+        # Test case 3: Multiple indices should create multiple sublists
+        result = split_list([1, 2, 3, 4, 5, 6, 7], [2, 5])
+        assert result == [[1, 2], [3, 4, 5], [6, 7]]
+
+        # Test case 4: Index at the beginning
+        result = split_list([1, 2, 3, 4], [0])
+        assert result == [[], [1, 2, 3, 4]]
+
+        # Test case 5: Index at the end
+        result = split_list([1, 2, 3, 4], [4])
+        assert result == [[1, 2, 3, 4], []]
+
+        # Test case 6: Empty list
+        result = split_list([], [])
+        assert result == [[]]
+
+        # Test case 7: String list to test with different data types
+        result = split_list(["a", "b", "c", "d"], [2])
+        assert result == [["a", "b"], ["c", "d"]]