brainstate 0.0.1__py2.py3-none-any.whl → 0.0.1.post20240622__py2.py3-none-any.whl

brainstate/nn/_rate_rnns.py

@@ -90,6 +90,9 @@ class ValinaRNNCell(RNNCell):
   def init_state(self, batch_size: int = None, **kwargs):
     self.h = ShortTermState(init.param(self._state_initializer, self.num_out, batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.h.value = init.param(self._state_initializer, self.num_out, batch_size)
+
   def update(self, x):
     xh = jnp.concatenate([x, self.h.value], axis=-1)
     h = self.W(xh)
@@ -147,6 +150,9 @@ class GRUCell(RNNCell):
   def init_state(self, batch_size: int = None, **kwargs):
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x):
     old_h = self.h.value
     xh = jnp.concatenate([x, old_h], axis=-1)
@@ -224,6 +230,9 @@ class MGUCell(RNNCell):
   def init_state(self, batch_size: int = None, **kwargs):
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x):
     old_h = self.h.value
     xh = jnp.concatenate([x, old_h], axis=-1)
@@ -327,6 +336,10 @@ class LSTMCell(RNNCell):
     self.c = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.c.value = init.param(self._state_initializer, [self.num_out], batch_size)
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x):
     h, c = self.h.value, self.c.value
     xh = jnp.concat([x, h], axis=-1)
@@ -379,6 +392,10 @@ class URLSTMCell(RNNCell):
     self.c = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
     self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
 
+  def reset_state(self, batch_size: int = None, **kwargs):
+    self.c.value = init.param(self._state_initializer, [self.num_out], batch_size)
+    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
+
   def update(self, x: ArrayLike) -> ArrayLike:
     h, c = self.h.value, self.c.value
     xh = jnp.concat([x, h], axis=-1)

brainstate/nn/_readout.py

@@ -66,6 +66,9 @@ class LeakyRateReadout(DnnLayer):
   def init_state(self, batch_size=None, **kwargs):
     self.r = ShortTermState(init.param(init.Constant(0.), self.out_size, batch_size))
 
+  def reset_state(self, batch_size=None, **kwargs):
+    self.r.value = init.param(init.Constant(0.), self.out_size, batch_size)
+
   def update(self, x):
     r = self.decay * self.r.value + x @ self.weight.value
     self.r.value = r
@@ -109,6 +112,9 @@ class LeakySpikeReadout(Neuron):
   def init_state(self, batch_size, **kwargs):
     self.V = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
 
+  def reset_state(self, batch_size, **kwargs):
+    self.V.value = init.param(init.Constant(0.), self.varshape, batch_size)
+
   @property
   def spike(self):
     return self.get_spike(self.V.value)

brainstate/optim/__init__.py

@@ -20,4 +20,3 @@ from ._sgd_optimizer import *
 from ._sgd_optimizer import __all__ as optimizer_all
 
 __all__ = scheduler_all + optimizer_all
-

brainstate/optim/_lr_scheduler_test.py

@@ -34,3 +34,16 @@ class TestMultiStepLR(unittest.TestCase):
         self.assertTrue(jnp.allclose(r, 0.001))
       else:
         self.assertTrue(jnp.allclose(r, 0.0001))
+
+  def test2(self):
+    lr = bst.transform.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
+    for i in range(40):
+      r = lr(i)
+      if i < 10:
+        self.assertEqual(r, 0.1)
+      elif i < 20:
+        self.assertTrue(jnp.allclose(r, 0.01))
+      elif i < 30:
+        self.assertTrue(jnp.allclose(r, 0.001))
+      else:
+        self.assertTrue(jnp.allclose(r, 0.0001))

brainstate/optim/_sgd_optimizer.py

@@ -18,11 +18,12 @@
 import functools
 from typing import Union, Dict, Optional, Tuple, Any, TypeVar
 
+import brainunit as bu
 import jax
 import jax.numpy as jnp
 
 from ._lr_scheduler import make_schedule, LearningRateScheduler
-from .. import environ, math
+from .. import environ
 from .._module import Module
 from .._state import State, LongTermState, StateDictManager, visible_state_dict
 
@@ -282,7 +283,7 @@ class Momentum(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.momentum_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.momentum_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -349,7 +350,7 @@ class MomentumNesterov(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.momentum_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.momentum_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -417,7 +418,7 @@ class Adagrad(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.cache_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.cache_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -500,8 +501,8 @@ class Adadelta(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.cache_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.delta_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.cache_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.delta_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     weight_values, grad_values, cache_values, delta_values = to_same_dict_tree(
@@ -574,7 +575,7 @@ class RMSProp(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.cache_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.cache_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -647,8 +648,8 @@ class Adam(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.m1_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.m2_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.m1_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.m2_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -730,7 +731,7 @@ class LARS(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.momentum_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.momentum_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -835,10 +836,10 @@ class Adan(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.exp_avg_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.exp_avg_sq_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.exp_avg_diff_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.pre_grad_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.exp_avg_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.exp_avg_sq_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.exp_avg_diff_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.pre_grad_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr = self.lr()
@@ -989,10 +990,10 @@ class AdamW(_WeightDecayOptimizer):
     for k, v in train_states.items():
       assert isinstance(v, State), f'"{k}" must be an instance of brainstate.State.'
       self.weight_states.add_unique_elem(k, v)
-      self.m1_states[k] = OptimState(math.tree_zeros_like(v.value))
-      self.m2_states[k] = OptimState(math.tree_zeros_like(v.value))
+      self.m1_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
+      self.m2_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
       if self.amsgrad:
-        self.vmax_states[k] = OptimState(math.tree_zeros_like(v.value))
+        self.vmax_states[k] = OptimState(bu.math.tree_zeros_like(v.value))
 
   def update(self, grads: dict):
     lr_old = self.lr()

brainstate/transform/__init__.py

@@ -17,10 +17,10 @@
 This module contains the functions for the transformation of the brain data.
 """
 
-from ._control import *
-from ._control import __all__ as _controls_all
 from ._autograd import *
 from ._autograd import __all__ as _gradients_all
+from ._control import *
+from ._control import __all__ as _controls_all
 from ._jit import *
 from ._jit import __all__ as _jit_all
 from ._jit_error import *
@@ -33,4 +33,3 @@ from ._progress_bar import __all__ as _progress_bar_all
 __all__ = _gradients_all + _jit_error_all + _controls_all + _make_jaxpr_all + _jit_all + _progress_bar_all
 
 del _gradients_all, _jit_error_all, _controls_all, _make_jaxpr_all, _jit_all, _progress_bar_all
-

brainstate/transform/_autograd.py

@@ -25,8 +25,8 @@ from jax._src.api import _vjp
 from jax.api_util import argnums_partial
 from jax.extend import linear_util
 
-from brainstate._utils import set_module_as
 from brainstate._state import State, StateTrace, StateDictManager
+from brainstate._utils import set_module_as
 
 __all__ = [
   'vector_grad', 'grad', 'jacrev', 'jacfwd', 'jacobian', 'hessian',

brainstate/transform/_autograd_test.py

@@ -537,7 +537,6 @@ class TestPureFuncJacobian(unittest.TestCase):
 
   def test_jacrev_return_aux1(self):
     with bc.environ.context(precision=64):
-
       def f1(x, y):
         a = 4 * x[1] ** 2 - 2 * x[2]
         r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], a, x[2] * jnp.sin(x[0])])
@@ -564,7 +563,6 @@ class TestPureFuncJacobian(unittest.TestCase):
       assert (vec == _r).all()
 
 
-
 class TestClassFuncJacobian(unittest.TestCase):
   def test_jacrev1(self):
     def f1(x, y):

brainstate/transform/_jit.py

@@ -23,8 +23,8 @@ import jax
 from jax._src import sharding_impls
 from jax.lib import xla_client as xc
 
-from ._make_jaxpr import StatefulFunction, _ensure_index_tuple, _assign_state_values
 from brainstate._utils import set_module_as
+from ._make_jaxpr import StatefulFunction, _ensure_index_tuple, _assign_state_values
 
 __all__ = ['jit']
 
@@ -33,10 +33,13 @@ class JittedFunction(Callable):
   """
   A wrapped version of ``fun``, set up for just-in-time compilation.
   """
-  origin_fun: Callable   # the original function
+  origin_fun: Callable  # the original function
   stateful_fun: StatefulFunction  # the stateful function for extracting states
   jitted_fun: jax.stages.Wrapped  # the jitted function
-  clear_cache: Callable # clear the cache of the jitted function
+  clear_cache: Callable  # clear the cache of the jitted function
+
+  def __call__(self, *args, **kwargs):
+    pass
 
 
 def _get_jitted_fun(
@@ -85,12 +88,16 @@ def _get_jitted_fun(
     jit_fun.clear_cache()
 
   jitted_fun: JittedFunction
+
   # the original function
   jitted_fun.origin_fun = fun.fun
+
   # the stateful function for extracting states
   jitted_fun.stateful_fun = fun
+
   # the jitted function
   jitted_fun.jitted_fun = jit_fun
+
   # clear cache
   jitted_fun.clear_cache = clear_cache
 
@@ -99,18 +106,18 @@ def _get_jitted_fun(
 
 @set_module_as('brainstate.transform')
 def jit(
-  fun: Callable = None,
-  in_shardings=sharding_impls.UNSPECIFIED,
-  out_shardings=sharding_impls.UNSPECIFIED,
-  static_argnums: int | Sequence[int] | None = None,
-  donate_argnums: int | Sequence[int] | None = None,
-  donate_argnames: str | Iterable[str] | None = None,
-  keep_unused: bool = False,
-  device: xc.Device | None = None,
-  backend: str | None = None,
-  inline: bool = False,
-  abstracted_axes: Any | None = None,
-  **kwargs
+    fun: Callable = None,
+    in_shardings=sharding_impls.UNSPECIFIED,
+    out_shardings=sharding_impls.UNSPECIFIED,
+    static_argnums: int | Sequence[int] | None = None,
+    donate_argnums: int | Sequence[int] | None = None,
+    donate_argnames: str | Iterable[str] | None = None,
+    keep_unused: bool = False,
+    device: xc.Device | None = None,
+    backend: str | None = None,
+    inline: bool = False,
+    abstracted_axes: Any | None = None,
+    **kwargs
 ) -> Union[JittedFunction, Callable[[Callable], JittedFunction]]:
   """
   Sets up ``fun`` for just-in-time compilation with XLA.
@@ -228,12 +235,31 @@ def jit(
 
   if fun is None:
     def wrapper(fun_again: Callable) -> JittedFunction:
-      return _get_jitted_fun(fun_again, in_shardings, out_shardings, static_argnums,
-                             donate_argnums, donate_argnames, keep_unused,
-                             device, backend, inline, abstracted_axes, **kwargs)
+      return _get_jitted_fun(fun_again,
+                             in_shardings,
+                             out_shardings,
+                             static_argnums,
+                             donate_argnums,
+                             donate_argnames,
+                             keep_unused,
+                             device,
+                             backend,
+                             inline,
+                             abstracted_axes,
+                             **kwargs)
+
     return wrapper
 
   else:
-    return _get_jitted_fun(fun, in_shardings, out_shardings, static_argnums,
-                           donate_argnums, donate_argnames, keep_unused,
-                           device, backend, inline, abstracted_axes, **kwargs)
+    return _get_jitted_fun(fun,
+                           in_shardings,
+                           out_shardings,
+                           static_argnums,
+                           donate_argnums,
+                           donate_argnames,
+                           keep_unused,
+                           device,
+                           backend,
+                           inline,
+                           abstracted_axes,
+                           **kwargs)

brainstate/transform/_jit_test.py

@@ -16,7 +16,6 @@
 import unittest
 
 import jax.numpy as jnp
-import jax.stages
 
 import brainstate as bc
 
@@ -90,7 +89,6 @@ class TestJIT(unittest.TestCase):
     self.assertTrue(len(compiling) == 2)
 
   def test_jit_attribute_origin_fun(self):
-
     def fun1(x):
       return x
 
@@ -99,4 +97,3 @@ class TestJIT(unittest.TestCase):
     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

@@ -54,20 +54,26 @@ function.
 from __future__ import annotations
 
 import functools
+import inspect
 import operator
 from collections.abc import Hashable, Iterable, Sequence
+from contextlib import ExitStack
 from typing import Any, Callable, Tuple, Union, Dict, Optional
 
 import jax
 from jax._src import source_info_util
+from jax._src.linear_util import annotate
+from jax._src.traceback_util import api_boundary
+from jax.extend.linear_util import transformation_with_aux, wrap_init
 from jax.interpreters import partial_eval as pe
-from jax.util import wraps
 from jax.interpreters.xla import abstractify
+from jax.util import wraps
 
 from brainstate._state import State, StateTrace
 from brainstate._utils import set_module_as
 
 PyTree = Any
+AxisName = Hashable
 
 __all__ = [
   "StatefulFunction",
@@ -393,7 +399,8 @@ class StatefulFunction(object):
     if cache_key not in self._state_trace:
       try:
         # jaxpr
-        jaxpr, (out_shapes, state_shapes) = jax.make_jaxpr(
+        # jaxpr, (out_shapes, state_shapes) = jax.make_jaxpr(
+        jaxpr, (out_shapes, state_shapes) = _make_jaxpr(
           functools.partial(self._wrapped_fun_to_eval, cache_key),
           static_argnums=self.static_argnums,
           axis_env=self.axis_env,
@@ -474,7 +481,8 @@ def make_jaxpr(
     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.
+  """
+  Creates a function that produces its jaxpr given example args.
 
   Args:
     fun: The function whose ``jaxpr`` is to be computed. Its positional
@@ -571,3 +579,156 @@ def make_jaxpr(
   if hasattr(fun, "__name__"):
     make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})"
   return make_jaxpr_f
+
+
+def _check_callable(fun):
+  # In Python 3.10+, the only thing stopping us from supporting staticmethods
+  # is that we can't take weak references to them, which the C++ JIT requires.
+  if isinstance(fun, staticmethod):
+    raise TypeError(f"staticmethod arguments are not supported, got {fun}")
+  if not callable(fun):
+    raise TypeError(f"Expected a callable value, got {fun}")
+  if inspect.isgeneratorfunction(fun):
+    raise TypeError(f"Expected a function, got a generator function: {fun}")
+
+
+def _broadcast_prefix(
+    prefix_tree: Any,
+    full_tree: Any,
+    is_leaf: Callable[[Any], bool] | None = None
+) -> list[Any]:
+  # If prefix_tree is not a tree prefix of full_tree, this code can raise a
+  # ValueError; use prefix_errors to find disagreements and raise more precise
+  # error messages.
+  result = []
+  num_leaves = lambda t: jax.tree.structure(t).num_leaves
+  add_leaves = lambda x, subtree: result.extend([x] * num_leaves(subtree))
+  jax.tree.map(add_leaves, prefix_tree, full_tree, is_leaf=is_leaf)
+  return result
+
+
+def _flat_axes_specs(
+    abstracted_axes, *args, **kwargs
+) -> list[pe.AbstractedAxesSpec]:
+  if kwargs:
+    raise NotImplementedError
+
+  def ax_leaf(l):
+    return (isinstance(l, dict) and jax.tree_util.all_leaves(l.values()) or
+            isinstance(l, tuple) and jax.tree_util.all_leaves(l, lambda x: x is None))
+
+  return _broadcast_prefix(abstracted_axes, args, ax_leaf)
+
+
+@transformation_with_aux
+def _flatten_fun(in_tree, *args_flat):
+  py_args, py_kwargs = jax.tree.unflatten(in_tree, args_flat)
+  ans = yield py_args, py_kwargs
+  yield jax.tree.flatten(ans)
+
+
+def _make_jaxpr(
+    fun: Callable,
+    static_argnums: int | Iterable[int] = (),
+    axis_env: Sequence[tuple[AxisName, int]] | None = None,
+    return_shape: bool = False,
+    abstracted_axes: Any | None = None,
+) -> Callable[..., (jax.core.ClosedJaxpr | tuple[jax.core.ClosedJaxpr, Any])]:
+  """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
+      ``ClosedJaxpr`` representation of ``fun`` 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.
+
+  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
+  >>>
+  >>> def f(x): return jax.numpy.sin(jax.numpy.cos(x))
+  >>> print(f(3.0))
+  -0.83602
+  >>> _make_jaxpr(f)(3.0)
+  { lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) }
+  >>> _make_jaxpr(jax.grad(f))(3.0)
+  { 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,) }
+  """
+  _check_callable(fun)
+  static_argnums = _ensure_index_tuple(static_argnums)
+
+  def _abstractify(args, kwargs):
+    flat_args, in_tree = jax.tree.flatten((args, kwargs))
+    if abstracted_axes is None:
+      return map(jax.api_util.shaped_abstractify, flat_args), in_tree, [True] * len(flat_args)
+    else:
+      axes_specs = _flat_axes_specs(abstracted_axes, *args, **kwargs)
+      in_type = pe.infer_lambda_input_type(axes_specs, flat_args)
+      in_avals, keep_inputs = jax.util.unzip2(in_type)
+      return in_avals, in_tree, keep_inputs
+
+  @wraps(fun)
+  @api_boundary
+  def make_jaxpr_f(*args, **kwargs):
+    f = wrap_init(fun)
+    if static_argnums:
+      dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
+      f, args = jax.api_util.argnums_partial(f, dyn_argnums, args)
+    in_avals, in_tree, keep_inputs = _abstractify(args, kwargs)
+    in_type = tuple(jax.util.safe_zip(in_avals, keep_inputs))
+    f, out_tree = _flatten_fun(f, in_tree)
+    f = annotate(f, in_type)
+    debug_info = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
+    with ExitStack() as stack:
+      for axis_name, size in axis_env or []:
+        stack.enter_context(jax.core.extend_axis_env(axis_name, size, None))
+      jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
+    closed_jaxpr = jax.core.ClosedJaxpr(jaxpr, consts)
+    if return_shape:
+      out_avals, _ = jax.util.unzip2(out_type)
+      out_shapes_flat = [jax.ShapeDtypeStruct(a.shape, a.dtype, a.named_shape) for a in out_avals]
+      return closed_jaxpr, jax.tree.unflatten(out_tree(), out_shapes_flat)
+    return closed_jaxpr
+
+  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

brainstate/transform/_make_jaxpr_test.py

@@ -129,5 +129,3 @@ def test_return_states():
 
   with pytest.raises(ValueError):
     f()
-
-

brainstate/transform/_progress_bar.py

@@ -14,13 +14,12 @@
 # ==============================================================================
 
 from __future__ import annotations
+
 import copy
 from typing import Optional
 
 import jax
 
-from brainstate import environ
-
 try:
   from tqdm.auto import tqdm
 except (ImportError, ModuleNotFoundError):
@@ -95,7 +94,6 @@ class ProgressBarRunner(object):
     self.tqdm_bars[0].close()
 
   def __call__(self, iter_num, *args, **kwargs):
-
     _ = jax.lax.cond(
       iter_num == 0,
       lambda: jax.debug.callback(self._define_tqdm),

brainstate/util.py

@@ -26,7 +26,6 @@ from jax.lib import xla_bridge
 
 from ._utils import set_module_as
 
-
 __all__ = [
   'unique_name',
   'clear_buffer_memory',

{brainstate-0.0.1.dist-info → brainstate-0.0.1.post20240622.dist-info}/METADATA

@@ -1,8 +1,8 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.1
+Version: 0.0.1.post20240622
 Summary: A State-based Transformation System for Brain Dynamics Programming.
-Home-page: https://github.com/brainpy/braincore
+Home-page: https://github.com/brainpy/brainstate
 Author: BrainPy Team
 Author-email: BrainPy Team <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -31,21 +31,11 @@ License-File: LICENSE
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
+Requires-Dist: brainunit
 Provides-Extra: cpu
 Requires-Dist: jaxlib ; extra == 'cpu'
-Requires-Dist: brainpylib ; extra == 'cpu'
-Provides-Extra: cpu_mini
-Requires-Dist: jaxlib ; extra == 'cpu_mini'
-Provides-Extra: cuda11
-Requires-Dist: jaxlib[cuda11_pip] ; extra == 'cuda11'
-Requires-Dist: brainpylib ; extra == 'cuda11'
-Provides-Extra: cuda11_mini
-Requires-Dist: jaxlib[cuda11_pip] ; extra == 'cuda11_mini'
 Provides-Extra: cuda12
 Requires-Dist: jaxlib[cuda12_pip] ; extra == 'cuda12'
-Requires-Dist: brainpylib ; extra == 'cuda12'
-Provides-Extra: cuda12_mini
-Requires-Dist: jaxlib[cuda12_pip] ; extra == 'cuda12_mini'
 Provides-Extra: testing
 Requires-Dist: pytest ; extra == 'testing'
 Provides-Extra: tpu
@@ -90,12 +80,14 @@ The official documentation is hosted on Read the Docs: [https://brainstate.readt
 
 ## See also the BDP ecosystem
 
-- [``brainpy``](https://github.com/brainpy/BrainPy): The solution for the general-purpose brain dynamics programming.
+- [``brainstate``](https://github.com/brainpy/brainstate): A ``State``-based transformation system for brain dynamics programming.
 
-- [``brainstate``](https://github.com/brainpy/brainstate): The core system for the next generation of BrainPy framework.
+- [``brainunit``](https://github.com/brainpy/brainunit): The unit system for brain dynamics programming.
 
-- [``braintools``](https://github.com/brainpy/braintools): The tools for the brain dynamics simulation and analysis.
+- [``braintaichi``](https://github.com/brainpy/braintaichi): Leveraging Taichi Lang to customize brain dynamics operators.
 
-- [``brainscale``](https://github.com/brainpy/brainscale): The scalable online learning for biological spiking neural networks.
+- [``brainscale``](https://github.com/brainpy/brainscale): The scalable online learning framework for biological neural networks.
+
+- [``braintools``](https://github.com/brainpy/braintools): The toolbox for the brain dynamics simulation, training and analysis.