Merge branch 'develop' into develop

Author: kratman

docs/source/api/util.rst

@@ -16,6 +16,4 @@ Utility functions
 
 .. autofunction:: pybamm.has_jax
 
-.. autofunction:: pybamm.is_jax_compatible
-
 .. autofunction:: pybamm.set_logging_level

pyproject.toml

@@ -79,8 +79,7 @@ plot = [
     "matplotlib>=3.6.0",
 ]
 cite = [
-    "setuptools", # Fix for a pybtex issue
-    "pybtex>=0.24.0",
+    "pybtex>=0.25.0",
 ]
 # Battery Parameter eXchange format
 bpx = [
@@ -112,13 +111,9 @@ dev = [
     "importlib-metadata; python_version < '3.10'",
     # For property based testing
     "hypothesis",
-
 ]
-# For the Jax solver.
-# Note: These must be kept in sync with the versions defined in pybamm/util.py
 jax = [
-    "jax==0.4.27",
-    "jaxlib==0.4.27",
+    "jax>=0.4.36,<0.6.0",
 ]
 # Contains all optional dependencies, except for jax, and dev dependencies
 all = [

src/pybamm/__init__.py

@@ -12,7 +12,6 @@
     get_parameters_filepath,
     has_jax,
     import_optional_dependency,
-    is_jax_compatible,
 )
 from .logger import logger, set_logging_level, get_new_logger
 from .settings import settings

src/pybamm/solvers/idaklu_jax.py

@@ -600,6 +600,7 @@ 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"
             )

src/pybamm/solvers/jax_bdf_solver.py

@@ -910,7 +910,7 @@ def merge(l1, l2):
         return out, merge
 
     def abstractify(x):
-        return core.raise_to_shaped(core.get_aval(x))
+        return core.get_aval(x)
 
     def ravel_first_arg(f, unravel):
         return ravel_first_arg_(lu.wrap_init(f), unravel).call_wrapped

src/pybamm/util.py

@@ -12,11 +12,6 @@
 
 import pybamm
 
-# Versions of jax and jaxlib compatible with PyBaMM. Note: these are also defined in
-# the extras dependencies in pyproject.toml, and therefore must be kept in sync.
-JAX_VERSION = "0.4.27"
-JAXLIB_VERSION = "0.4.27"
-
 
 def root_dir():
     """return the root directory of the PyBaMM install directory"""
@@ -354,27 +349,11 @@ def has_jax():
         True if jax and jaxlib are installed with the correct versions, False if otherwise
 
     """
-    return (
-        (importlib.util.find_spec("jax") is not None)
-        and (importlib.util.find_spec("jaxlib") is not None)
-        and is_jax_compatible()
+    return (importlib.util.find_spec("jax") is not None) and (
+        importlib.util.find_spec("jaxlib") is not None
     )
 
 
-def is_jax_compatible():
-    """
-    Check if the available versions of jax and jaxlib are compatible with PyBaMM
-
-    Returns
-    -------
-    bool
-        True if jax and jaxlib are compatible with PyBaMM, False if otherwise
-    """
-    return importlib.metadata.distribution("jax").version.startswith(
-        JAX_VERSION
-    ) and importlib.metadata.distribution("jaxlib").version.startswith(JAXLIB_VERSION)
-
-
 def is_constant_and_can_evaluate(symbol):
     """
     Returns True if symbol is constant and evaluation does not raise any errors.

tests/unit/test_util.py

@@ -90,10 +90,6 @@ def test_get_parameters_filepath(self, tmp_path):
             os.path.join(pybamm.root_dir(), "src", "pybamm", temppath)
         )
 
-    @pytest.mark.skipif(not pybamm.has_jax(), reason="JAX is not installed")
-    def test_is_jax_compatible(self):
-        assert pybamm.is_jax_compatible()
-
     def test_import_optional_dependency(self):
         optional_distribution_deps = get_optional_distribution_deps("pybamm")
         present_optional_import_deps = get_present_optional_import_deps(