brainstate 0.1.3__py2.py3-none-any.whl → 0.1.5__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.3"
+__version__ = "0.1.5"
 
 from . import augment
 from . import compile

brainstate/_compatible_import.py

@@ -16,10 +16,9 @@
 # -*- coding: utf-8 -*-
 
 
-import importlib.util
 from contextlib import contextmanager
 from functools import partial
-from typing import Iterable, Hashable, TypeVar, Callable, TYPE_CHECKING
+from typing import Iterable, Hashable, TypeVar, Callable
 
 import jax
 
@@ -31,7 +30,6 @@ __all__ = [
     'get_aval',
     'Tracer',
     'to_concrete_aval',
-    'brainevent',
     'safe_map',
     'safe_zip',
     'unzip2',
@@ -47,8 +45,6 @@ T3 = TypeVar("T3")
 
 from saiunit._compatible_import import wrap_init
 
-brainevent_installed = importlib.util.find_spec('brainevent') is not None
-
 from jax.core import get_aval, Tracer
 
 if jax.__version_info__ < (0, 5, 0):
@@ -150,14 +146,3 @@ def to_concrete_aval(aval):
         return aval.to_concrete_value()
     return aval
 
-
-if not brainevent_installed:
-    if not TYPE_CHECKING:
-        class BrainEvent:
-            def __getattr__(self, item):
-                raise ImportError('brainevent is not installed, please install brainevent first.')
-
-        brainevent = BrainEvent()
-
-else:
-    import brainevent

brainstate/_state.py

@@ -50,6 +50,7 @@ __all__ = [
     'LongTermState',
     'HiddenState',
     'ParamState',
+    'BatchState',
     'TreefyState',
     'FakeState',
 

brainstate/augment/_mapping.py

@@ -185,10 +185,10 @@ def _compile_stateful_function(
     if isinstance(in_axes, int):
         args = jax.tree.map(lambda x: _remove_axis(x, in_axes), args)
     elif isinstance(in_axes, tuple):
-        args = tuple(
-            [arg if in_axis is None else _remove_axis(arg, in_axis)
-             for arg, in_axis in zip(args, in_axes)]
-        )
+        args = tuple([
+            arg if in_axis is None else _remove_axis(arg, in_axis)
+            for arg, in_axis in zip(args, in_axes)
+        ])
     stateful_fn.make_jaxpr(state_vals, args)
     return stateful_fn.get_arg_cache_key(state_vals, args)
 
@@ -383,10 +383,7 @@ def _vmap_transform(
         stateful_fn.axis_env = axis_env
 
     # stateful function
-    stateful_fn = StatefulFunction(
-        _vmap_fn_for_compilation,
-        name='vmap',
-    )
+    stateful_fn = StatefulFunction(_vmap_fn_for_compilation, name='vmap')
 
     @functools.wraps(f)
     def new_fn_for_vmap(
@@ -460,7 +457,10 @@ def _vmap_transform(
         # analyze vmapping axis error
         for state in state_trace.get_write_states():
             leaves = jax.tree.leaves(state.value)
-            if any([isinstance(leaf, BatchTracer) for leaf in leaves]) and state not in out_state_to_axis:
+            if (
+                any([isinstance(leaf, BatchTracer) and (leaf.batch_dim is not None) for leaf in leaves])
+                and state not in out_state_to_axis
+            ):
                 if isinstance(state, RandomState) and state in rng_sets:
                     continue
                 state.raise_error_with_source_info(

brainstate/augment/_mapping_test.py

@@ -19,6 +19,8 @@ import unittest
 import jax
 import jax.numpy as jnp
 import numpy as np
+from jax import vmap
+from jax.lax import psum, pmean, pmax
 
 import brainstate
 import brainstate.augment
@@ -433,3 +435,163 @@ class TestVMAPNewStatesEdgeCases(unittest.TestCase):
             foo.c = brainstate.State(jnp.arange(3))  # Original expected shape is (4,)
 
         faulty_init()
+
+
+class TestAxisName:
+    def test1(self):
+        def compute_stats_with_axis_name(x):
+            """Compute statistics using named axis operations"""
+            # Sum across the named axis 'batch'
+            total_sum = psum(x, axis_name='batch')
+
+            # Mean across the named axis 'batch'
+            mean_val = pmean(x, axis_name='batch')
+
+            # Max across the named axis 'batch'
+            max_val = pmax(x, axis_name='batch')
+
+            return {
+                'sum': total_sum,
+                'mean': mean_val,
+                'max': max_val,
+                'original': x
+            }
+
+        batch_data = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
+        print("Input batch data:", batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats_jax = jax.jit(vmap(compute_stats_with_axis_name, axis_name='batch'))
+        result_jax = vectorized_stats_jax(batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats = brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
+        result = vectorized_stats(batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats_v2 = brainstate.transform.jit(
+            brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
+        )
+        result_v2 = vectorized_stats_v2(batch_data)
+
+        for key in result_jax.keys():
+            print(f"  {key}: {result_jax[key]}")
+            assert jnp.allclose(result_jax[key], result[key]), f"Mismatch in {key}"
+            assert jnp.allclose(result_jax[key], result_v2[key]), f"Mismatch in {key}"
+
+    def test_nested_vmap(self):
+        def nested_computation(x):
+            """Computation with multiple named axes"""
+            # Sum over 'inner' axis, then mean over 'outer' axis
+            inner_sum = psum(x, axis_name='inner')
+            outer_mean = pmean(inner_sum, axis_name='outer')
+            return outer_mean
+
+        # Create 2D batch data
+        data_2d = jnp.arange(12.0).reshape(3, 4)  # Shape: [outer_batch=3, inner_batch=4]
+        print("Input 2D data shape:", data_2d.shape)
+        print("Input 2D data:\n", data_2d)
+
+        # Nested vmap: first over inner dimension, then outer dimension
+        inner_vmap = vmap(nested_computation, axis_name='inner')
+        nested_vmap = vmap(inner_vmap, axis_name='outer')
+
+        result_2d = nested_vmap(data_2d)
+        print("Result after nested vmap:", result_2d)
+
+        inner_vmap_bst = brainstate.transform.vmap(nested_computation, axis_name='inner')
+        nested_vmap_bst = brainstate.transform.vmap(inner_vmap_bst, axis_name='outer')
+        result_2d_bst = nested_vmap_bst(data_2d)
+        print("Result after nested vmap:", result_2d_bst)
+
+        assert jnp.allclose(result_2d, result_2d_bst)
+
+    def _gradient_averaging_simulation_bst(self):
+        def loss_function(params, x, y):
+            """Simple quadratic loss"""
+            pred = params * x
+            return (pred - y) ** 2
+
+        def compute_gradients_with_averaging(params, batch_x, batch_y):
+            """Compute gradients and average them across the batch"""
+            # Compute per-sample gradients
+            grad_fn = jax.grad(loss_function, argnums=0)
+            per_sample_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+
+            # Average gradients across batch using named axis
+            def average_grads(grads):
+                return pmean(grads, axis_name='batch')
+
+            # Apply averaging with named axis
+            averaged_grads = vmap(average_grads, axis_name='batch')(per_sample_grads)
+            return averaged_grads
+
+        # Example data
+        params = 2.0
+        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
+        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
+
+        print("Parameters:", params)
+        print("Batch X:", batch_x)
+        print("Batch Y:", batch_y)
+
+        # Compute individual gradients first
+        grad_fn = jax.grad(loss_function, argnums=0)
+        individual_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+        print("Individual gradients:", individual_grads)
+
+        # Now compute averaged gradients using axis names
+        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
+        print("Averaged gradients:", averaged_grads)
+
+        return individual_grads, averaged_grads
+
+    def _gradient_averaging_simulation_jax(self):
+        def loss_function(params, x, y):
+            """Simple quadratic loss"""
+            pred = params * x
+            return (pred - y) ** 2
+
+        def compute_gradients_with_averaging(params, batch_x, batch_y):
+            """Compute gradients and average them across the batch"""
+            # Compute per-sample gradients
+            grad_fn = jax.grad(loss_function, argnums=0)
+            per_sample_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+
+            # Average gradients across batch using named axis
+            def average_grads(grads):
+                return pmean(grads, axis_name='batch')
+
+            # Apply averaging with named axis
+            averaged_grads = brainstate.transform.vmap(average_grads, axis_name='batch')(per_sample_grads)
+            return averaged_grads
+
+        # Example data
+        params = 2.0
+        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
+        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
+
+        print("Parameters:", params)
+        print("Batch X:", batch_x)
+        print("Batch Y:", batch_y)
+
+        # Compute individual gradients first
+        grad_fn = jax.grad(loss_function, argnums=0)
+        individual_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+        print("Individual gradients:", individual_grads)
+
+        # Now compute averaged gradients using axis names
+        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
+        print("Averaged gradients:", averaged_grads)
+
+        return individual_grads, averaged_grads
+
+    def test_gradient_averaging_simulation(self):
+        individual_grads, averaged_grads = self._gradient_averaging_simulation_bst()
+        individual_grads_jax, averaged_grads_jax = self._gradient_averaging_simulation_jax()
+        assert jnp.allclose(individual_grads, individual_grads_jax)
+        assert jnp.allclose(averaged_grads, averaged_grads_jax)
+
+
+
+

brainstate/compile/_jit.py

@@ -51,6 +51,7 @@ def _get_jitted_fun(
     out_shardings,
     static_argnums,
     donate_argnums,
+    static_argnames,
     donate_argnames,
     keep_unused,
     device,
@@ -59,10 +60,12 @@ def _get_jitted_fun(
     abstracted_axes,
     **kwargs
 ) -> JittedFunction:
-    static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
+    static_argnums = tuple() if static_argnums is None else _ensure_index_tuple(static_argnums)
+    donate_argnums = tuple() if donate_argnums is None else _ensure_index_tuple(donate_argnums)
     fun = StatefulFunction(
         fun,
         static_argnums=static_argnums,
+        static_argnames=static_argnames,
         abstracted_axes=abstracted_axes,
         cache_type='jit',
         name='jit'
@@ -70,7 +73,8 @@ def _get_jitted_fun(
     jit_fun = jax.jit(
         fun.jaxpr_call,
         static_argnums=tuple(i + 1 for i in static_argnums),
-        donate_argnums=donate_argnums,
+        static_argnames=static_argnames,
+        donate_argnums=tuple(i + 1 for i in donate_argnums),
         donate_argnames=donate_argnames,
         keep_unused=keep_unused,
         device=device,
@@ -179,6 +183,7 @@ def jit(
     out_shardings=sharding_impls.UNSPECIFIED,
     static_argnums: int | Sequence[int] | None = None,
     donate_argnums: int | Sequence[int] | None = None,
+    static_argnames: str | Sequence[str] | None = None,
     donate_argnames: str | Iterable[str] | None = None,
     keep_unused: bool = False,
     device: Device | None = None,
@@ -190,9 +195,6 @@ def jit(
     """
     Sets up ``fun`` for just-in-time compilation with XLA.
 
-    Does not support setting ``static_argnames`` as in ``jax.jit()``.
-
-
     Args:
       fun: Function to be jitted.
       in_shardings: Pytree of structure matching that of arguments to ``fun``,
@@ -246,6 +248,11 @@ def jit(
         provided, ``inspect.signature`` is not used, and only actual
         parameters listed in either ``static_argnums`` or ``static_argnames`` will
         be treated as static.
+      static_argnames: An optional string or collection of strings specifying
+        which named arguments are treated as static (compile-time constant).
+        Operations that only depend on static arguments will be constant-folded in
+        Python (during tracing), and so the corresponding argument values can be
+        any Python object.
       donate_argnums: Specify which positional argument buffers are "donated" to
         the computation. It is safe to donate argument buffers if you no longer
         need them once the computation has finished. In some cases XLA can make
@@ -309,6 +316,7 @@ def jit(
                 out_shardings=out_shardings,
                 static_argnums=static_argnums,
                 donate_argnums=donate_argnums,
+                static_argnames=static_argnames,
                 donate_argnames=donate_argnames,
                 keep_unused=keep_unused,
                 device=device,
@@ -327,6 +335,7 @@ def jit(
             out_shardings,
             static_argnums,
             donate_argnums,
+            static_argnames,
             donate_argnames,
             keep_unused,
             device,

brainstate/compile/_make_jaxpr.py

@@ -88,6 +88,12 @@ __all__ = [
 ]
 
 
+def _ensure_str(x: str) -> str:
+    if not isinstance(x, str):
+        raise TypeError(f"argument is not a string: {x}")
+    return x
+
+
 def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
     """Convert x to a tuple of indices."""
     x = jax.core.concrete_or_error(None, x, "expected a static index or sequence of indices.")
@@ -97,6 +103,14 @@ def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
         return tuple(safe_map(operator.index, x))
 
 
+def _ensure_str_tuple(x: str | Iterable[str]) -> tuple[str, ...]:
+    """Convert x to a tuple of strings."""
+    if isinstance(x, str):
+        return (x,)
+    else:
+        return tuple(safe_map(_ensure_str, x))
+
+
 def _jax_v04_new_arg_fn(frame, trace, aval):
     """
     Transform a new argument to a tracer.
@@ -155,6 +169,9 @@ def _init_state_trace_stack(name) -> StateTraceStack:
     return state_trace
 
 
+default_cache_key = ((), ())
+
+
 class StatefulFunction(PrettyObject):
     """
     A wrapper class for a function that collects the states that are read and written by the function. The states are
@@ -170,6 +187,7 @@ class StatefulFunction(PrettyObject):
         arguments and return value should be arrays, scalars, or standard Python
         containers (tuple/list/dict) thereof.
       static_argnums: See the :py:func:`jax.jit` docstring.
+      static_argnames: See the :py:func:`jax.jit` docstring.
       axis_env: Optional, a sequence of pairs where the first element is an axis
           name and the second element is a positive integer representing the size of
           the mapped axis with that name. This parameter is useful when lowering
@@ -199,6 +217,7 @@ class StatefulFunction(PrettyObject):
         self,
         fun: Callable,
         static_argnums: Union[int, Iterable[int]] = (),
+        static_argnames: Union[str, Iterable[str]] = (),
         axis_env: Optional[Sequence[tuple[Hashable, int]]] = None,
         abstracted_axes: Optional[Any] = None,
         state_returns: Union[str, Tuple[str, ...]] = ('read', 'write'),
@@ -207,11 +226,12 @@ class StatefulFunction(PrettyObject):
     ):
         # explicit parameters
         self.fun = fun
-        self.static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
+        self.static_argnums = tuple() if static_argnums is None else _ensure_index_tuple(static_argnums)
+        self.static_argnames = tuple() if static_argnames is None else _ensure_str_tuple(static_argnames)
         self.axis_env = axis_env
         self.abstracted_axes = abstracted_axes
         self.state_returns = tuple(state_returns) if isinstance(state_returns, (tuple, list)) else (state_returns,)
-        assert cache_type in [None, 'jit']
+        assert cache_type in [None, 'jit'], f"Invalid cache type: {cache_type}"
         self.name = name
 
         # implicit parameters
@@ -226,7 +246,7 @@ class StatefulFunction(PrettyObject):
             return None
         return k, v
 
-    def get_jaxpr(self, cache_key: Hashable = ()) -> ClosedJaxpr:
+    def get_jaxpr(self, cache_key: Hashable = None) -> ClosedJaxpr:
         """
         Read the JAX Jaxpr representation of the function.
 
@@ -236,11 +256,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The JAX Jaxpr representation of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_jaxpr:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_jaxpr[cache_key]
 
-    def get_out_shapes(self, cache_key: Hashable = ()) -> PyTree:
+    def get_out_shapes(self, cache_key: Hashable = None) -> PyTree:
         """
         Read the output shapes of the function.
 
@@ -250,11 +272,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The output shapes of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_out_shapes:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_out_shapes[cache_key]
 
-    def get_out_treedef(self, cache_key: Hashable = ()) -> PyTree:
+    def get_out_treedef(self, cache_key: Hashable = None) -> PyTree:
         """
         Read the output tree of the function.
 
@@ -264,11 +288,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The output tree of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_jaxpr_out_tree:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_jaxpr_out_tree[cache_key]
 
-    def get_state_trace(self, cache_key: Hashable = ()) -> StateTraceStack:
+    def get_state_trace(self, cache_key: Hashable = None) -> StateTraceStack:
         """
         Read the state trace of the function.
 
@@ -278,11 +304,13 @@ class StatefulFunction(PrettyObject):
         Returns:
           The state trace of the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         if cache_key not in self._cached_state_trace:
             raise ValueError(f"the function is not called with the static arguments: {cache_key}")
         return self._cached_state_trace[cache_key]
 
-    def get_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    def get_states(self, cache_key: Hashable = None) -> Tuple[State, ...]:
         """
         Read the states that are read and written by the function.
 
@@ -292,9 +320,11 @@ class StatefulFunction(PrettyObject):
         Returns:
           The states that are read and written by the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         return tuple(self.get_state_trace(cache_key).states)
 
-    def get_read_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    def get_read_states(self, cache_key: Hashable = None) -> Tuple[State, ...]:
         """
         Read the states that are read by the function.
 
@@ -304,9 +334,11 @@ class StatefulFunction(PrettyObject):
         Returns:
           The states that are read by the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         return self.get_state_trace(cache_key).get_read_states()
 
-    def get_write_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    def get_write_states(self, cache_key: Hashable = None) -> Tuple[State, ...]:
         """
         Read the states that are written by the function.
 
@@ -316,6 +348,8 @@ class StatefulFunction(PrettyObject):
         Returns:
           The states that are written by the function.
         """
+        if cache_key is None:
+            cache_key = default_cache_key
         return self.get_state_trace(cache_key).get_write_states()
 
     def get_arg_cache_key(self, *args, **kwargs) -> Tuple:
@@ -323,10 +357,11 @@ class StatefulFunction(PrettyObject):
         Get the static arguments from the arguments.
 
         Args:
-          *args: The arguments to the function.
+            *args: The arguments to the function.
+            **kwargs: The keyword arguments to the function.
 
         Returns:
-          The static arguments.
+          The static arguments and keyword arguments as a tuple.
         """
         if self.cache_type == 'jit':
             static_args, dyn_args = [], []
@@ -336,11 +371,18 @@ class StatefulFunction(PrettyObject):
                 else:
                     dyn_args.append(arg)
             dyn_args = jax.tree.map(shaped_abstractify, jax.tree.leaves(dyn_args))
-            dyn_kwargs = jax.tree.map(shaped_abstractify, jax.tree.leaves(kwargs))
-            return tuple([tuple(static_args), tuple(dyn_args), tuple(dyn_kwargs)])
+            static_kwargs, dyn_kwargs = [], []
+            for k, v in kwargs.items():
+                if k in self.static_argnames:
+                    static_kwargs.append((k, v))
+                else:
+                    dyn_kwargs.append((k, jax.tree.map(shaped_abstractify, v)))
+            return tuple([tuple(static_args), tuple(dyn_args), tuple(static_kwargs), tuple(dyn_kwargs)])
         elif self.cache_type is None:
             num_arg = len(args)
-            return tuple(args[i] for i in self.static_argnums if i < num_arg)
+            static_args = tuple(args[i] for i in self.static_argnums if i < num_arg)
+            static_kwargs = tuple((k, v) for k, v in kwargs.items() if k in self.static_argnames)
+            return tuple([static_args, static_kwargs])
         else:
             raise ValueError(f"Invalid cache type: {self.cache_type}")
 
@@ -389,7 +431,7 @@ class StatefulFunction(PrettyObject):
         self._cached_state_trace.clear()
 
     def _wrapped_fun_to_eval(
-        self, cache_key, *args, return_only_write: bool = False, **kwargs,
+        self, cache_key, static_kwargs: dict, *args, return_only_write: bool = False, **dyn_kwargs,
     ) -> Tuple[Any, Tuple[State, ...]]:
         """
         Wrap the function and return the states that are read and written by the function and the output of the function.
@@ -405,7 +447,7 @@ class StatefulFunction(PrettyObject):
         state_trace = _init_state_trace_stack(self.name)
         self._cached_state_trace[cache_key] = state_trace
         with state_trace:
-            out = self.fun(*args, **kwargs)
+            out = self.fun(*args, **dyn_kwargs, **static_kwargs)
             state_values = (
                 state_trace.get_write_state_values(True)
                 if return_only_write else
@@ -430,8 +472,9 @@ class StatefulFunction(PrettyObject):
         the structure, shape, dtypes, and named shapes of the output of ``fun``.
 
         Args:
-          *args: The arguments to the function.
-          **kwargs: The keyword arguments to the function.
+            *args: The arguments to the function.
+            **kwargs: The keyword arguments to the function.
+            return_only_write: If True, only return the states that are written by the function.
         """
 
         # static args
@@ -440,17 +483,24 @@ class StatefulFunction(PrettyObject):
         if cache_key not in self._cached_state_trace:
             try:
                 # jaxpr
+                static_kwargs, dyn_kwargs = {}, {}
+                for k, v in kwargs.items():
+                    if k in self.static_argnames:
+                        static_kwargs[k] = v
+                    else:
+                        dyn_kwargs[k] = v
                 jaxpr, (out_shapes, state_shapes) = _make_jaxpr(
                     functools.partial(
                         self._wrapped_fun_to_eval,
                         cache_key,
+                        static_kwargs,
                         return_only_write=return_only_write
                     ),
                     static_argnums=self.static_argnums,
                     axis_env=self.axis_env,
                     return_shape=True,
                     abstracted_axes=self.abstracted_axes
-                )(*args, **kwargs)
+                )(*args, **dyn_kwargs)
                 # returns
                 self._cached_jaxpr_out_tree[cache_key] = jax.tree.structure((out_shapes, state_shapes))
                 self._cached_out_shapes[cache_key] = (out_shapes, state_shapes)
@@ -483,6 +533,7 @@ class StatefulFunction(PrettyObject):
         assert len(state_vals) == len(states), 'State length mismatch.'
 
         # parameters
+        kwargs = {k: v for k, v in kwargs.items() if k not in self.static_argnames}  # remove static kwargs
         args = tuple(args[i] for i in range(len(args)) if i not in self.static_argnums)
         args = jax.tree.flatten((args, kwargs, state_vals))[0]
 
@@ -519,12 +570,16 @@ class StatefulFunction(PrettyObject):
 def make_jaxpr(
     fun: Callable,
     static_argnums: Union[int, Iterable[int]] = (),
+    static_argnames: Union[str, Iterable[str]] = (),
     axis_env: Optional[Sequence[tuple[Hashable, int]]] = None,
     return_shape: bool = False,
     abstracted_axes: Optional[Any] = None,
     state_returns: Union[str, Tuple[str, ...]] = ('read', 'write')
-) -> Callable[..., (Tuple[ClosedJaxpr, Tuple[State, ...]] |
-                    Tuple[ClosedJaxpr, Tuple[State, ...], PyTree])]:
+) -> Callable[
+    ...,
+    (Tuple[ClosedJaxpr, Tuple[State, ...]] |
+     Tuple[ClosedJaxpr, Tuple[State, ...], PyTree])
+]:
     """
     Creates a function that produces its jaxpr given example args.
 
@@ -533,6 +588,7 @@ def make_jaxpr(
         arguments and return value should be arrays, scalars, or standard Python
         containers (tuple/list/dict) thereof.
       static_argnums: See the :py:func:`jax.jit` docstring.
+      static_argnames: See the :py:func:`jax.jit` docstring.
       axis_env: Optional, a sequence of pairs where the first element is an axis
         name and the second element is a positive integer representing the size of
         the mapped axis with that name. This parameter is useful when lowering
@@ -605,11 +661,11 @@ def make_jaxpr(
     stateful_fun = StatefulFunction(
         fun,
         static_argnums=static_argnums,
+        static_argnames=static_argnames,
         axis_env=axis_env,
         abstracted_axes=abstracted_axes,
         state_returns=state_returns,
         name='make_jaxpr'
-
     )
 
     @wraps(fun)