brainstate 0.0.1__py2.py3-none-any.whl

brainstate/transform/_jit_test.py

@@ -0,0 +1,102 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import unittest
+
+import jax.numpy as jnp
+import jax.stages
+
+import brainstate as bc
+
+
+class TestJIT(unittest.TestCase):
+  def test_inner_state_are_not_catched(self):
+    a = bc.State(bc.random.randn(10))
+
+    @bc.transform.jit
+    def fun1(inp):
+      a.value += inp
+
+      b = bc.State(bc.random.randn(1))
+
+      def inner_fun(x):
+        b.value += x
+
+      bc.transform.for_loop(inner_fun, bc.random.randn(100))
+
+      return a.value + b.value
+
+    print(fun1(1.))
+    key = fun1.stateful_fun.get_arg_cache_key(1.)
+    self.assertTrue(len(fun1.stateful_fun.get_states(key)) == 2)
+
+    x = bc.random.randn(10)
+    print(fun1(x))
+    key = fun1.stateful_fun.get_arg_cache_key(x)
+    self.assertTrue(len(fun1.stateful_fun.get_states(key)) == 2)
+
+  def test_jit_compile_sensitive_to_input_shape(self):
+    global_data = [0]
+
+    @bc.transform.jit
+    def fun1(inp):
+      global_data[0] += 1
+      return inp
+
+    print(fun1(1.))
+    self.assertTrue(global_data[0] == 1)
+
+    print(fun1(2.))
+    self.assertTrue(global_data[0] == 1)
+
+    print(fun1(bc.random.randn(10)))
+    self.assertTrue(global_data[0] == 2)
+
+    print(fun1(bc.random.randn(10, 10)))
+    self.assertTrue(global_data[0] == 3)
+
+  def test_jit_clear_cache(self):
+    a = bc.State(bc.random.randn(1))
+    compiling = []
+
+    @bc.transform.jit
+    def log2(x):
+      print('compiling')
+      compiling.append(1)
+      ln_x = jnp.log(x)
+      ln_2 = jnp.log(2.0) + a.value
+      return ln_x / ln_2
+
+    x = bc.random.randn(1)
+    print(log2(x))  # compiling
+    self.assertTrue(len(compiling) == 1)
+    print(log2(x))  # no compiling
+    self.assertTrue(len(compiling) == 1)
+
+    log2.clear_cache()
+    print(log2(x))  # compiling
+    self.assertTrue(len(compiling) == 2)
+
+  def test_jit_attribute_origin_fun(self):
+
+    def fun1(x):
+      return x
+
+    jitted_fun = bc.transform.jit(fun1)
+    self.assertTrue(jitted_fun.origin_fun is fun1)
+    self.assertTrue(isinstance(jitted_fun.stateful_fun, bc.transform.StatefulFunction))
+    self.assertTrue(callable(jitted_fun.jitted_fun))
+    self.assertTrue(callable(jitted_fun.clear_cache))
+

brainstate/transform/_make_jaxpr.py

@@ -0,0 +1,573 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""
+This module implements how to create a JAX Jaxpr from a given function by considering the states that are read and
+written by the function. These state transformations are foundational for the BrainCore library. These utilities
+include two basic functions: `StatefulFunction` and `make_jaxpr`.
+
+
+``StatefulFunction``
+--------------------
+
+The module provides a class called ``StatefulFunction`` that wraps a function and provides methods to get the
+JAX Jaxpr, the output shapes, the states that are read and written by the function, and the output of the function.
+The class provides the following methods:
+
+- `make_jaxpr`: creates the JAX Jaxpr of the function.
+- `jaxpr_call`: calls the function at the JAX Jaxpr level.
+- `jaxpr_call_without_states`: calls the function at the JAX Jaxpr level without considering the states.
+- `get_states`: returns the states that are read and written by the function.
+- `get_read_states`: returns the states that are read by the function.
+- `get_write_states`: returns the states that are written by the function.
+- `get_static_args`: returns the static arguments from the arguments.
+- `compile_and_get_states_by_static_args`: compiles the function and returns the states that are read and
+   written by the function.
+- `get_jaxpr`: returns the JAX Jaxpr of the function.
+- `get_out_shapes`: returns the output shapes of the function.
+- `get_out_treedef`: returns the output tree of the function.
+
+``make_jaxpr``
+--------------
+
+The module provides a function called `make_jaxpr` that creates a function that produces its JAX Jaxpr given example
+arguments. The function returns a wrapped version of the function that when applied to example arguments returns a
+`ClosedJaxpr` representation of the function on those arguments. If the argument `return_shape` is `True`, then the
+returned function instead returns a pair where the first element is the `ClosedJaxpr` representation of the function
+and the second element is a pytree representing the structure, shape, dtypes, and named shapes of the output of the
+function.
+
+"""
+
+from __future__ import annotations
+
+import functools
+import operator
+from collections.abc import Hashable, Iterable, Sequence
+from typing import Any, Callable, Tuple, Union, Dict, Optional
+
+import jax
+from jax._src import source_info_util
+from jax.interpreters import partial_eval as pe
+from jax.util import wraps
+from jax.interpreters.xla import abstractify
+
+from brainstate._state import State, StateTrace
+from brainstate._utils import set_module_as
+
+PyTree = Any
+
+__all__ = [
+  "StatefulFunction",
+  "make_jaxpr",
+]
+
+
+def _assign_state_values(states, state_vals) -> None:
+  """
+  Assign the state values to the states.
+
+  Args:
+    states: The states.
+    state_vals: The state values.
+  """
+  assert len(states) == len(state_vals), f'State length mismatch. {len(states)} != {len(state_vals)}.'
+  for st, val in zip(states, state_vals):
+    st.value = val
+
+
+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.")
+  try:
+    return (operator.index(x),)
+  except TypeError:
+    return tuple(jax.util.safe_map(operator.index, x))
+
+
+def _new_arg(frame, trace, aval):
+  """
+  Transform a new argument to a tracer.
+
+  Modified from jax.interpreters.partial_eval.DynamicJaxprTrace.new_arg()
+
+  Args:
+    frame: The frame.
+    trace: The trace.
+    aval: The abstract value.
+
+  Returns:
+    The tracer.
+  """
+  tracer = pe.DynamicJaxprTracer(trace, aval, source_info_util.current())
+  frame.tracers.append(tracer)
+  frame.tracer_to_var[id(tracer)] = var = frame.newvar(aval)
+  frame.invars.append(var)
+  return tracer
+
+
+def wrapped_abstractify(x: Any) -> Any:
+  """
+  Abstractify the input.
+
+  Args:
+    x: The input.
+
+  Returns:
+    The abstractified input.
+  """
+  if isinstance(x, pe.DynamicJaxprTracer):
+    return jax.core.ShapedArray(x.aval.shape, x.aval.dtype, weak_type=x.aval.weak_type)
+  return abstractify(x)
+
+
+class StatefulFunction(object):
+  """
+  A wrapper class for a function that collects the states that are read and written by the function. The states are
+  collected by the function and returned as a StateDictManager instance. The StateDictManager instance can be used to
+  manage the states in the JAX program. The class provides a function called `states` that returns the states
+  that are read and written by the function. The class provides a function called `to_state_manager` that returns
+  a StateDictManager instance that contains the states that are read and written by the function. The class provides
+  a function called `__call__` that wraps the function and returns the states that are read and written by the
+  function and the output of the function.
+
+  Args:
+    fun: The function whose ``jaxpr`` is to be computed. Its positional
+      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.
+    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
+        functions that involve parallel communication collectives, and it
+        specifies the axis name/size environment that would be set up by
+        applications of :py:func:`jax.pmap`.
+    abstracted_axes: Optional, a pytree with the same structure as the input
+        arguments to ``fun``. The leaves of the pytree can be either None or a
+        dict with axis names as keys and integers as values. If the leaf is None,
+        then the corresponding axis is not abstracted. If the leaf is a dict, then
+        the corresponding axis is abstracted, and the dict specifies the axis name
+        and size. The abstracted axes are used to infer the input type of the
+        function. If None, then all axes are abstracted.
+    state_returns: Optional, a string or a tuple of strings. The default is
+        ``('read', 'write')``. The strings specify the categories of states to be
+        returned by the wrapped function. The categories are ``'read'`` and
+        ``'write'``. If the category is ``'read'``, then the wrapped function
+        returns the states that are read by the function. If the category is
+        ``'write'``, then the wrapped function returns the states that are written
+        by the function. If the category is ``'read'`` and ``'write'``, then the
+        wrapped function returns both the read and write states.
+
+  """
+  __module__ = "brainstate.transform"
+
+  def __init__(
+      self,
+      fun: Callable,
+      static_argnums: Union[int, Iterable[int]] = (),
+      axis_env: Optional[Sequence[tuple[Hashable, int]]] = None,
+      abstracted_axes: Optional[Any] = None,
+      state_returns: Union[str, Tuple[str, ...]] = ('read', 'write'),
+      cache_type: Optional[str] = None,
+  ):
+    # explicit parameters
+    self.fun = fun
+    self.static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
+    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']
+    self.cache_type = cache_type
+
+    # implicit parameters
+    self._jaxpr: Dict[Any, jax.core.ClosedJaxpr] = dict()
+    self._out_shapes: Dict[Any, PyTree] = dict()
+    self._jaxpr_out_tree: Dict[Any, PyTree] = dict()
+    self._state_trace: Dict[Any, StateTrace] = dict()
+
+  def __repr__(self) -> str:
+    return (f"{self.__class__.__name__}({self.fun}, "
+            f"static_argnums={self.static_argnums}, "
+            f"axis_env={self.axis_env}, "
+            f"abstracted_axes={self.abstracted_axes}, "
+            f"state_returns={self.state_returns})")
+
+  def get_jaxpr(self, cache_key: Hashable = ()) -> jax.core.ClosedJaxpr:
+    """
+    Read the JAX Jaxpr representation of the function.
+
+    Args:
+      cache_key: The hashable key.
+
+    Returns:
+      The JAX Jaxpr representation of the function.
+    """
+    if cache_key not in self._jaxpr:
+      raise ValueError(f"the function is not called with the static arguments: {cache_key}")
+    return self._jaxpr[cache_key]
+
+  def get_out_shapes(self, cache_key: Hashable = ()) -> PyTree:
+    """
+    Read the output shapes of the function.
+
+    Args:
+      cache_key: The hashable key.
+
+    Returns:
+      The output shapes of the function.
+    """
+    if cache_key not in self._out_shapes:
+      raise ValueError(f"the function is not called with the static arguments: {cache_key}")
+    return self._out_shapes[cache_key]
+
+  def get_out_treedef(self, cache_key: Hashable = ()) -> PyTree:
+    """
+    Read the output tree of the function.
+
+    Args:
+      cache_key: The hashable key.
+
+    Returns:
+      The output tree of the function.
+    """
+    if cache_key not in self._jaxpr_out_tree:
+      raise ValueError(f"the function is not called with the static arguments: {cache_key}")
+    return self._jaxpr_out_tree[cache_key]
+
+  def get_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    """
+    Read the states that are read and written by the function.
+
+    Args:
+      cache_key: The hashable key.
+
+    Returns:
+      The states that are read and written by the function.
+    """
+    if cache_key not in self._state_trace:
+      raise ValueError(f"the function is not called with the static arguments: {cache_key}")
+    return tuple(self._state_trace[cache_key].states)
+
+  def get_read_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    """
+    Read the states that are read by the function.
+
+    Args:
+      cache_key: The hashable key.
+
+    Returns:
+      The states that are read by the function.
+    """
+    _state_trace = self._state_trace[cache_key]
+    return tuple([st for st, ty in zip(_state_trace.states, _state_trace.types) if ty == 'read'])
+
+  def get_write_states(self, cache_key: Hashable = ()) -> Tuple[State, ...]:
+    """
+    Read the states that are written by the function.
+
+    Args:
+      cache_key: The hashable key.
+
+    Returns:
+      The states that are written by the function.
+    """
+    state_trace = self._state_trace[cache_key]
+    return tuple([st for st, ty in zip(state_trace.states, state_trace.types) if ty == 'write'])
+
+  def get_arg_cache_key(self, *args, **kwargs) -> Tuple:
+    """
+    Get the static arguments from the arguments.
+
+    Args:
+      *args: The arguments to the function.
+
+    Returns:
+      The static arguments.
+    """
+    if self.cache_type == 'jit':
+      static_args, dyn_args = [], []
+      for i, arg in enumerate(args):
+        if i in self.static_argnums:
+          static_args.append(arg)
+        else:
+          dyn_args.append(arg)
+      dyn_args = jax.tree.map(wrapped_abstractify, jax.tree.leaves(dyn_args))
+      dyn_kwargs = jax.tree.map(wrapped_abstractify, jax.tree.leaves(kwargs))
+      return tuple([tuple(static_args), tuple(dyn_args), 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)
+    else:
+      raise ValueError(f"Invalid cache type: {self.cache_type}")
+
+  def compile_and_get_states_by_static_args(self, *args, **kwargs) -> Tuple[State, ...]:
+    """
+    Get the states that are read and written by the function.
+
+    Args:
+      *args: The arguments to the function.
+      **kwargs: The keyword arguments to the function.
+
+    Returns:
+      The states that are read and written by the function.
+    """
+    cache_key = self.get_arg_cache_key(*args, **kwargs)
+    if cache_key not in self._state_trace:
+      self.make_jaxpr(*args, **kwargs)
+    return self.get_states(cache_key)
+
+  def clear_cache(self) -> None:
+    """
+    Clear the compilation cache.
+    """
+    self._jaxpr.clear()
+    self._out_shapes.clear()
+    self._jaxpr_out_tree.clear()
+    self._state_trace.clear()
+
+  @staticmethod
+  def _init_trace_and_newarg() -> StateTrace:
+    # Should be within the calling of ``jax.make_jaxpr()``
+    state_trace: StateTrace = StateTrace()
+    main = jax.core.thread_local_state.trace_state.trace_stack.stack[-1]
+    frame = main.jaxpr_stack[-1]
+    trace = pe.DynamicJaxprTrace(main, jax.core.cur_sublevel())
+    state_trace.set_new_arg(functools.partial(_new_arg, frame, trace))
+    return state_trace
+
+  def _wrapped_fun_to_eval(self, cache_key, *args, **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.
+
+    Args:
+      *args: The arguments to the function.
+      **kwargs: The keyword arguments to the function.
+
+    Returns:
+      A tuple of the states that are read and written by the function and the output of the function.
+    """
+    # state trace
+    _state_trace = self._init_trace_and_newarg()
+    self._state_trace[cache_key] = _state_trace
+    with _state_trace:
+      out = self.fun(*args, **kwargs)
+      state_values = _state_trace.collect_values('read', 'write')
+    _state_trace.recovery_original_values()
+
+    # return states is not allowed
+    # checking whether the states are returned
+    for leaf in jax.tree.leaves(out):
+      if isinstance(leaf, State):
+        leaf._raise_error_with_source_info(ValueError(f"State object is not allowed to be returned: {leaf}"))
+    return out, state_values
+
+  def make_jaxpr(self, *args, **kwargs):
+    """Creates a function that produces its jaxpr given example args.
+
+    A ``ClosedJaxpr`` representation of ``fun`` on those arguments. If the
+    argument ``return_shape`` is ``True``, then the returned function instead
+    returns a pair where the first element is the ``ClosedJaxpr``
+    representation of ``fun`` and the second element is a pytree representing
+    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.
+    """
+
+    # static args
+    cache_key = self.get_arg_cache_key(*args, **kwargs)
+
+    if cache_key not in self._state_trace:
+      try:
+        # jaxpr
+        jaxpr, (out_shapes, state_shapes) = jax.make_jaxpr(
+          functools.partial(self._wrapped_fun_to_eval, cache_key),
+          static_argnums=self.static_argnums,
+          axis_env=self.axis_env,
+          return_shape=True,
+          abstracted_axes=self.abstracted_axes
+        )(*args, **kwargs)
+
+        # returns
+        self._jaxpr_out_tree[cache_key] = jax.tree.structure((out_shapes, state_shapes))
+        self._out_shapes[cache_key] = (out_shapes, state_shapes)
+        self._jaxpr[cache_key] = jaxpr
+      except Exception as e:
+        try:
+          self._state_trace.pop(cache_key)
+        except KeyError:
+          pass
+        raise e
+
+    return self
+
+  def jaxpr_call(self, state_vals, *args, **kwargs) -> Any:
+    """
+    Call the function at the JAX Jaxpr level.
+
+    Args:
+      state_vals: The state values.
+      *args: The arguments to the function.
+      **kwargs: The keyword arguments to the function.
+
+    Returns:
+      State values and the function output.
+    """
+    # state checking
+    cache_key = self.get_arg_cache_key(*args, **kwargs)
+    states = self.get_states(cache_key)
+    assert len(state_vals) == len(states), 'State length mismatch.'
+
+    # parameters
+    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]
+
+    # calling the function
+    closed_jaxpr = self.get_jaxpr(cache_key)
+    out_treedef = self.get_out_treedef(cache_key)
+    jaxpr_outs = jax.core.eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.consts, *args)
+
+    # output processing
+    out, new_state_vals = out_treedef.unflatten(jaxpr_outs)
+    assert len(new_state_vals) == len(state_vals), 'State length mismatch.'
+    return new_state_vals, out
+
+  def jaxpr_call_auto(self, *args, **kwargs) -> Any:
+    """
+    Call the function at the JAX Jaxpr level with automatic state management.
+
+    Args:
+      *args: The arguments to the function.
+      **kwargs: The keyword arguments to the function.
+
+    Returns:
+      The output of the function.
+    """
+    cache_key = self.get_arg_cache_key(*args, **kwargs)
+    states = self.get_states(cache_key)
+    state_vals, out = self.jaxpr_call([st.value for st in states], *args, **kwargs)
+    for st, val in zip(states, state_vals):
+      st.value = val
+    return out
+
+
+@set_module_as("brainstate.transform")
+def make_jaxpr(
+    fun: Callable,
+    static_argnums: Union[int, Iterable[int]] = (),
+    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[jax.core.ClosedJaxpr, Tuple[State, ...]] |
+                    Tuple[jax.core.ClosedJaxpr, Tuple[State, ...], PyTree])]:
+  """Creates a function that produces its jaxpr given example args.
+
+  Args:
+    fun: The function whose ``jaxpr`` is to be computed. Its positional
+      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.
+    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
+      functions that involve parallel communication collectives, and it
+      specifies the axis name/size environment that would be set up by
+      applications of :py:func:`jax.pmap`.
+    return_shape: Optional boolean, defaults to ``False``. If ``True``, the
+      wrapped function returns a pair where the first element is the XLA
+      computation and the second element is a pytree with the same structure as
+      the output of ``fun`` and where the leaves are objects with ``shape``,
+      ``dtype``, and ``named_shape`` attributes representing the corresponding
+      types of the output leaves.
+    abstracted_axes: Optional, a pytree with the same structure as the input
+      arguments to ``fun``. The leaves of the pytree can be either None or a
+      dict with axis names as keys and integers as values. If the leaf is None,
+      then the corresponding axis is not abstracted. If the leaf is a dict, then
+      the corresponding axis is abstracted, and the dict specifies the axis name
+      and size. The abstracted axes are used to infer the input type of the
+      function. If None, then all axes are abstracted.
+    state_returns: Optional, a string or a tuple of strings. The default is
+      ``('read', 'write')``. The strings specify the categories of states to be
+      returned by the wrapped function. The categories are ``'read'`` and
+      ``'write'``. If the category is ``'read'``, then the wrapped function
+      returns the states that are read by the function. If the category is
+      ``'write'``, then the wrapped function returns the states that are written
+      by the function. If the category is ``'read'`` and ``'write'``, then the
+      wrapped function returns both the read and write states.
+
+
+  Returns:
+    A wrapped version of ``fun`` that when applied to example arguments returns
+    a ``ClosedJaxpr`` representation of ``fun`` on those arguments. If the
+    argument ``return_shape`` is ``True``, then the returned function instead
+    returns a pair where the first element is the ``ClosedJaxpr``
+    representation of ``fun`` and the second element is a pytree representing
+    the structure, shape, dtypes, and named shapes of the output of ``fun``.
+
+  A ``jaxpr`` is JAX's intermediate representation for program traces. The
+  ``jaxpr`` language is based on the simply-typed first-order lambda calculus
+  with let-bindings. :py:func:`make_jaxpr` adapts a function to return its
+  ``jaxpr``, which we can inspect to understand what JAX is doing internally.
+  The ``jaxpr`` returned is a trace of ``fun`` abstracted to
+  :py:class:`ShapedArray` level. Other levels of abstraction exist internally.
+
+  We do not describe the semantics of the ``jaxpr`` language in detail here, but
+  instead give a few examples.
+
+  >>> import jax
+  >>> import brainstate as bst
+  >>>
+  >>> def f(x): return jax.numpy.sin(jax.numpy.cos(x))
+  >>> print(f(3.0))
+  -0.83602
+  >>> jaxpr, states = bst.transform.make_jaxpr(f)(3.0)
+  >>> jaxpr
+  { lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) }
+  >>> jaxpr, states = bst.transform.make_jaxpr(jax.grad(f))(3.0)
+  >>> jaxpr
+  { lambda ; a:f32[]. let
+      b:f32[] = cos a
+      c:f32[] = sin a
+      _:f32[] = sin b
+      d:f32[] = cos b
+      e:f32[] = mul 1.0 d
+      f:f32[] = neg e
+      g:f32[] = mul f c
+    in (g,) }
+  """
+
+  stateful_fun = StatefulFunction(fun, static_argnums, axis_env, abstracted_axes, state_returns)
+
+  @wraps(fun)
+  def make_jaxpr_f(*args, **kwargs):
+    stateful_fun.make_jaxpr(*args, **kwargs)
+    cache_key = stateful_fun.get_arg_cache_key(*args, **kwargs)
+    if return_shape:
+      return (stateful_fun.get_jaxpr(cache_key),
+              stateful_fun.get_states(cache_key),
+              stateful_fun.get_out_shapes(cache_key)[0])
+    else:
+      return (stateful_fun.get_jaxpr(cache_key),
+              stateful_fun.get_states(cache_key))
+
+  # wrapped jaxpr builder function
+  make_jaxpr_f.__module__ = "brainstate.transform"
+  if hasattr(fun, "__qualname__"):
+    make_jaxpr_f.__qualname__ = f"make_jaxpr({fun.__qualname__})"
+  if hasattr(fun, "__name__"):
+    make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})"
+  return make_jaxpr_f