Simplify JAX compat: use jax.make_jaxpr and aval helpers (#137)

* Refactor import statements for clarity and update jax version checks; enhance error callback functionality

* Refactor to use compatible import for get_aval in _loop_collect_return.py

* Add mapped_aval import and update references in loop_collect_return.py

* Refactor _compatible_import.py and _make_jaxpr.py: remove unused imports and functions for clarity

* chore: update jax dependency version to >=0.6.0 in pyproject.toml and requirements.txt

* refactor(test): remove test_all_exports for clarity and maintainability

* refactor(test): remove test_function_imports_availability for clarity

Author: chaoming0625

brainstate/_compatible_import.py

@@ -42,9 +42,8 @@
     >>> # These imports work across different JAX versions
 """
 
-from contextlib import contextmanager
 from functools import partial
-from typing import Iterable, Hashable, TypeVar, Callable
+from typing import Iterable, TypeVar, Callable
 
 import jax
 from jax.core import Tracer
@@ -76,9 +75,9 @@
     # others
     'is_jit_primitive',
     'Primitive',
-    'extend_axis_env_nd',
     'jaxpr_as_fun',
     'get_aval',
+    'mapped_aval',
     'to_concrete_aval',
     'Device',
     'wrap_init',
@@ -105,12 +104,15 @@ def get_aval(x):
 else:
     from jax import Device
 
+if jax.__version_info__ < (0, 8, 2):
+    from jax.core import mapped_aval
+else:
+    from jax.extend.core import mapped_aval
 if jax.__version_info__ < (0, 8, 0):
     from jax.lib.xla_bridge import get_backend
 else:
     from jax.extend.backend import get_backend
 
-
 if jax.__version_info__ < (0, 7, 1):
     from jax.interpreters.batching import make_iota, to_elt, BatchTracer, BatchTrace
 else:
@@ -119,40 +121,15 @@ def get_aval(x):
 from jax.core import DropVar
 
 if jax.__version_info__ < (0, 4, 38):
-    from jax.core import ClosedJaxpr, extend_axis_env_nd, Primitive, jaxpr_as_fun
-    from jax.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
+    from jax.core import (
+        extend_axis_env_nd, jaxpr_as_fun,
+        ClosedJaxpr, Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
+    )
 else:
-    from jax.extend.core import ClosedJaxpr, Primitive, jaxpr_as_fun
-    from jax.extend.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
-    from jax.core import trace_ctx
-
-
-    @contextmanager
-    def extend_axis_env_nd(name_size_pairs: Iterable[tuple[Hashable, int]]):
-        """
-        Context manager to temporarily extend the JAX axis environment.
-
-        Extends the current JAX axis environment with new named axes for
-        vectorized computations, then restores the previous environment.
-
-        Args:
-            name_size_pairs: Iterable of (name, size) tuples specifying
-                           the named axes to add to the environment.
-
-        Yields:
-            None: Context with extended axis environment.
-
-        Examples:
-            >>> with extend_axis_env_nd([('batch', 32), ('seq', 128)]):
-            ...     # Code using vectorized operations with named axes
-            ...     pass
-        """
-        prev = trace_ctx.axis_env
-        try:
-            trace_ctx.set_axis_env(prev.extend_pure(name_size_pairs))
-            yield
-        finally:
-            trace_ctx.set_axis_env(prev)
+    from jax.extend.core import (
+        ClosedJaxpr, jaxpr_as_fun,
+        Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
+    )
 
 if jax.__version_info__ < (0, 6, 0):
     from jax.util import safe_map, safe_zip, unzip2, wraps
@@ -257,30 +234,6 @@ def unzip2(xys: Iterable[tuple[T1, T2]]) -> tuple[tuple[T1, ...], tuple[T2, ...]
 
 
     def fun_name(fun: Callable):
-        """
-        Extract the name of a function, handling special cases.
-
-        Attempts to get the name of a function, with special handling for
-        partial functions and fallback for unnamed functions.
-
-        Args:
-            fun: The function to get the name from.
-
-        Returns:
-            str: The function name, or "<unnamed function>" if no name available.
-
-        Examples:
-            >>> def my_function():
-            ...     pass
-            >>> fun_name(my_function)
-            'my_function'
-
-            >>> from functools import partial
-            >>> add = lambda x, y: x + y
-            >>> add_one = partial(add, 1)
-            >>> fun_name(add_one)
-            '<lambda>'
-        """
         name = getattr(fun, "__name__", None)
         if name is not None:
             return name

brainstate/_compatible_import_test.py

@@ -71,34 +71,6 @@ def test_core_imports_availability(self):
                             f"{type_name} should be available")
             self.assertIsNotNone(getattr(compat, type_name))
 
-    def test_function_imports_availability(self):
-        """Test function imports are available."""
-        functions = [
-            'jaxpr_as_fun', 'get_aval', 'to_concrete_aval',
-            'extend_axis_env_nd'
-        ]
-
-        for func_name in functions:
-            self.assertTrue(hasattr(compat, func_name),
-                            f"{func_name} should be available")
-            self.assertTrue(callable(getattr(compat, func_name)),
-                            f"{func_name} should be callable")
-
-    def test_extend_axis_env_nd_functionality(self):
-        """Test extend_axis_env_nd context manager."""
-        # Test basic functionality
-        with compat.extend_axis_env_nd([('test_axis', 10)]):
-            # Context should execute without error
-            pass
-
-        # Test with multiple axes
-        with compat.extend_axis_env_nd([('batch', 32), ('seq', 128)]):
-            pass
-
-        # Test with empty axes
-        with compat.extend_axis_env_nd([]):
-            pass
-
     def test_get_aval_functionality(self):
         """Test get_aval function works correctly."""
         # Test with JAX array
@@ -650,21 +622,6 @@ def test_function(x):
 class TestModuleStructure(unittest.TestCase):
     """Test module structure and __all__ exports."""
 
-    def test_all_exports(self):
-        """Test that __all__ contains expected exports."""
-        expected_exports = [
-            'ClosedJaxpr', 'Primitive', 'extend_axis_env_nd', 'jaxpr_as_fun',
-            'get_aval', 'Tracer', 'to_concrete_aval', 'safe_map', 'safe_zip',
-            'unzip2', 'wraps', 'Device', 'wrap_init', 'Var', 'JaxprEqn',
-            'Jaxpr', 'Literal'
-        ]
-
-        for export in expected_exports:
-            self.assertIn(export, compat.__all__,
-                          f"{export} should be in __all__")
-            self.assertTrue(hasattr(compat, export),
-                            f"{export} should be available in module")
-
     def test_no_unexpected_exports(self):
         """Test that no private functions are exported."""
         for name in compat.__all__:

brainstate/transform/_error_if.py

@@ -28,7 +28,7 @@
 
 
 def _err_jit_true_branch(err_fun, args, kwargs):
-    jax.debug.callback(err_fun, *args, **kwargs)
+    jax.debug.callback(err_fun, *args, **kwargs, ordered=True)
 
 
 def _err_jit_false_branch(args, kwargs):

brainstate/transform/_jit.py

@@ -15,7 +15,7 @@
 
 import functools
 from collections.abc import Iterable, Sequence
-from typing import (Any, Callable, Union)
+from typing import Callable, Union
 
 import jax
 from jax._src import sharding_impls

brainstate/transform/_loop_collect_return.py

@@ -20,6 +20,7 @@
 import jax
 import jax.numpy as jnp
 
+from brainstate._compatible_import import get_aval, mapped_aval
 from brainstate._utils import set_module_as
 from ._make_jaxpr import StatefulFunction
 from ._progress_bar import ProgressBar
@@ -242,8 +243,8 @@ def scan(
 
     # evaluate jaxpr, get all states #
     # ------------------------------ #
-    xs_avals = [jax.core.get_aval(x) for x in xs_flat]
-    x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
+    xs_avals = [get_aval(x) for x in xs_flat]
+    x_avals = [mapped_aval(length, 0, aval) for aval in xs_avals]
     args = [init, xs_tree.unflatten(x_avals)]
     stateful_fun = StatefulFunction(f, name='scan').make_jaxpr(*args)
     state_trace = stateful_fun.get_state_trace(*args)
@@ -381,8 +382,8 @@ def checkpointed_scan(
         pbar_runner = None
 
     # evaluate jaxpr
-    xs_avals = [jax.core.get_aval(x) for x in xs_flat]
-    x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
+    xs_avals = [get_aval(x) for x in xs_flat]
+    x_avals = [mapped_aval(length, 0, aval) for aval in xs_avals]
     args = (init, xs_tree.unflatten(x_avals))
     stateful_fun = StatefulFunction(f, name='checkpoint_scan').make_jaxpr(*args)
     state_trace = stateful_fun.get_state_trace(*args)