brainstate 0.0.1.post20240623__py2.py3-none-any.whl → 0.0.1.1.post20240708__py2.py3-none-any.whl

brainstate/nn/_projection/_align_post.py

@@ -22,7 +22,7 @@ from brainstate._module import (register_delay_of_target,
                                 ReceiveInputProj,
                                 ExtendedUpdateWithBA)
 from brainstate._utils import set_module_as
-from brainstate.mixin import (Mode, AllOfTypes, DelayedInitializer, BindCondData, AlignPost, UpdateReturn)
+from brainstate.mixin import (Mode, JointTypes, DelayedInitializer, BindCondData, AlignPost, UpdateReturn)
 from ._utils import is_instance
 
 __all__ = [
@@ -39,7 +39,7 @@ def align_post_add_bef_update(
     out_label: str,
     syn_desc,
     out_desc,
-    post: AllOfTypes[ReceiveInputProj, ExtendedUpdateWithBA],
+    post: JointTypes[ReceiveInputProj, ExtendedUpdateWithBA],
     proj_name: str
 ):
   # synapse and output initialization
@@ -60,7 +60,7 @@ class _AlignPost(Module):
   def __init__(
       self,
       syn: Module,
-      out: AllOfTypes[Dynamics, BindCondData]
+      out: JointTypes[Dynamics, BindCondData]
   ):
     super().__init__()
     self.syn = syn
@@ -140,7 +140,7 @@ class HalfProjAlignPostMg(Projection):
       comm: Module,
       syn: DelayedInitializer[AlignPost],
       out: DelayedInitializer[BindCondData],
-      post: AllOfTypes[ReceiveInputProj, ExtendedUpdateWithBA],
+      post: JointTypes[ReceiveInputProj, ExtendedUpdateWithBA],
       out_label: Optional[str] = None,
       name: Optional[str] = None,
       mode: Optional[Mode] = None,
@@ -150,7 +150,7 @@ class HalfProjAlignPostMg(Projection):
     # synaptic models
     is_instance(syn, DelayedInitializer[AlignPost])
     is_instance(out, DelayedInitializer[BindCondData])
-    is_instance(post, AllOfTypes[ReceiveInputProj, ExtendedUpdateWithBA])
+    is_instance(post, JointTypes[ReceiveInputProj, ExtendedUpdateWithBA])
 
     # synapse and output initialization
     syn, out = align_post_add_bef_update(out_label, syn_desc=syn, out_desc=out, post=post, proj_name=self.name)
@@ -257,12 +257,12 @@ class FullProjAlignPostMg(Projection):
 
   def __init__(
       self,
-      pre: AllOfTypes[ExtendedUpdateWithBA, UpdateReturn],
+      pre: JointTypes[ExtendedUpdateWithBA, UpdateReturn],
       delay: Union[None, int, float],
       comm: Module,
       syn: DelayedInitializer[AlignPost],
       out: DelayedInitializer[BindCondData],
-      post: AllOfTypes[ReceiveInputProj, ExtendedUpdateWithBA],
+      post: JointTypes[ReceiveInputProj, ExtendedUpdateWithBA],
       out_label: Optional[str] = None,
       name: Optional[str] = None,
       mode: Optional[Mode] = None,
@@ -270,11 +270,11 @@ class FullProjAlignPostMg(Projection):
     super().__init__(name=name, mode=mode)
 
     # synaptic models
-    is_instance(pre, AllOfTypes[ExtendedUpdateWithBA, UpdateReturn])
+    is_instance(pre, JointTypes[ExtendedUpdateWithBA, UpdateReturn])
     is_instance(comm, Module)
     is_instance(syn, DelayedInitializer[AlignPost])
     is_instance(out, DelayedInitializer[BindCondData])
-    is_instance(post, AllOfTypes[ReceiveInputProj, ExtendedUpdateWithBA])
+    is_instance(post, JointTypes[ReceiveInputProj, ExtendedUpdateWithBA])
     self.comm = comm
 
     # delay initialization
@@ -494,8 +494,8 @@ class FullProjAlignPost(Projection):
     # synaptic models
     is_instance(pre, UpdateReturn)
     is_instance(comm, Module)
-    is_instance(syn, AllOfTypes[Dynamics, AlignPost])
-    is_instance(out, AllOfTypes[Dynamics, BindCondData])
+    is_instance(syn, JointTypes[Dynamics, AlignPost])
+    is_instance(out, JointTypes[Dynamics, BindCondData])
     is_instance(post, ReceiveInputProj)
     self.comm = comm
     self.syn = syn

brainstate/nn/_projection/_align_pre.py

@@ -20,7 +20,7 @@ from brainstate._module import (Module, DelayAccess, Projection,
                                 ExtendedUpdateWithBA, ReceiveInputProj,
                                 register_delay_of_target)
 from brainstate._utils import set_module_as
-from brainstate.mixin import (DelayedInitializer, BindCondData, UpdateReturn, Mode, AllOfTypes)
+from brainstate.mixin import (DelayedInitializer, BindCondData, UpdateReturn, Mode, JointTypes)
 from ._utils import is_instance
 
 __all__ = [
@@ -284,7 +284,7 @@ class FullProjAlignPreDSMg(Projection):
 
   def __init__(
       self,
-      pre: AllOfTypes[ExtendedUpdateWithBA, UpdateReturn],
+      pre: JointTypes[ExtendedUpdateWithBA, UpdateReturn],
       delay: Union[None, int, float],
       syn: DelayedInitializer[UpdateReturn],
       comm: Module,
@@ -297,7 +297,7 @@ class FullProjAlignPreDSMg(Projection):
     super().__init__(name=name, mode=mode)
 
     # synaptic models
-    is_instance(pre, AllOfTypes[ExtendedUpdateWithBA, UpdateReturn])
+    is_instance(pre, JointTypes[ExtendedUpdateWithBA, UpdateReturn])
     is_instance(syn, DelayedInitializer[Module])
     is_instance(comm, Module)
     is_instance(out, BindCondData)

brainstate/random.py

@@ -22,11 +22,16 @@ from functools import partial
 from operator import index
 from typing import Optional
 
+import brainunit as bu
 import jax
+import jax.numpy as jnp
+import jax.random as jr
 import numpy as np
-from jax import lax, jit, vmap, numpy as jnp, random as jr, core, dtypes
+from jax import jit, vmap
+from jax import lax, core, dtypes
 
 from brainstate import environ
+from ._random_for_unit import uniform_for_unit, permutation_for_unit
 from ._state import State
 from .transform._jit_error import jit_error
 from .typing import DTypeLike, Size, SeedOrKey
@@ -144,7 +149,7 @@ class RandomState(State):
   def rand(self, *dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
     key = self.split_key() if key is None else _formalize_key(key)
     dtype = dtype or environ.dftype()
-    r = jr.uniform(key, shape=dn, minval=0., maxval=1., dtype=dtype)
+    r = uniform_for_unit(key, shape=dn, minval=0., maxval=1., dtype=dtype)
     return r
 
   def randint(
@@ -207,7 +212,7 @@ class RandomState(State):
              dtype: DTypeLike = None):
     dtype = dtype or environ.dftype()
     key = self.split_key() if key is None else _formalize_key(key)
-    r = jr.uniform(key, shape=_size2shape(size), minval=0., maxval=1., dtype=dtype)
+    r = uniform_for_unit(key, shape=_size2shape(size), minval=0., maxval=1., dtype=dtype)
     return r
 
   def random_sample(self,
@@ -250,7 +255,7 @@ class RandomState(State):
                   key: Optional[SeedOrKey] = None):
     x = _check_py_seq(x)
     key = self.split_key() if key is None else _formalize_key(key)
-    r = jr.permutation(key, x, axis=axis, independent=independent)
+    r = permutation_for_unit(key, x, axis=axis, independent=independent)
     return r
 
   def shuffle(self,
@@ -258,7 +263,7 @@ class RandomState(State):
               axis=0,
               key: Optional[SeedOrKey] = None):
     key = self.split_key() if key is None else _formalize_key(key)
-    x = jr.permutation(key, x, axis=axis)
+    x = permutation_for_unit(key, x, axis=axis)
     return x
 
   def beta(self,
@@ -458,7 +463,7 @@ class RandomState(State):
       size = lax.broadcast_shapes(jnp.shape(low), jnp.shape(high))
     key = self.split_key() if key is None else _formalize_key(key)
     dtype = dtype or environ.dftype()
-    r = jr.uniform(key, shape=_size2shape(size), minval=low, maxval=high, dtype=dtype)
+    r = uniform_for_unit(key, shape=_size2shape(size), minval=low, maxval=high, dtype=dtype)
     return r
 
   def __norm_cdf(self, x, sqrt2, dtype):
@@ -481,36 +486,46 @@ class RandomState(State):
     scale = _check_py_seq(scale)
     dtype = dtype or environ.dftype()
 
-    lower = lax.convert_element_type(lower, dtype)
-    upper = lax.convert_element_type(upper, dtype)
-    loc = lax.convert_element_type(loc, dtype)
-    scale = lax.convert_element_type(scale, dtype)
+    lower = bu.math.asarray(lower, dtype=dtype)
+    upper = bu.math.asarray(upper, dtype=dtype)
+    loc = bu.math.asarray(loc, dtype=dtype)
+    scale = bu.math.asarray(scale, dtype=dtype)
+    bu.fail_for_dimension_mismatch(lower, upper)
+    bu.fail_for_dimension_mismatch(lower, loc)
+    bu.fail_for_dimension_mismatch(lower, scale)
+    dim = lower.dim if isinstance(lower, bu.Quantity) else bu.DIMENSIONLESS
+    lower = lower.value if isinstance(lower, bu.Quantity) else lower
+    upper = upper.value if isinstance(upper, bu.Quantity) else upper
+    loc = loc.value if isinstance(loc, bu.Quantity) else loc
+    scale = scale.value if isinstance(scale, bu.Quantity) else scale
 
     jit_error(
-      jnp.any(jnp.logical_or(loc < lower - 2 * scale, loc > upper + 2 * scale)),
+      bu.math.any(bu.math.logical_or(loc < lower - 2 * scale, loc > upper + 2 * scale)),
       "mean is more than 2 std from [lower, upper] in truncated_normal. "
       "The distribution of values may be incorrect."
     )
 
     if size is None:
-      size = lax.broadcast_shapes(jnp.shape(lower),
-                                  jnp.shape(upper),
-                                  jnp.shape(loc),
-                                  jnp.shape(scale))
+      size = bu.math.broadcast_shapes(jnp.shape(lower),
+                                      jnp.shape(upper),
+                                      jnp.shape(loc),
+                                      jnp.shape(scale))
 
     # Values are generated by using a truncated uniform distribution and
     # then using the inverse CDF for the normal distribution.
     # Get upper and lower cdf values
-    sqrt2 = np.array(np.sqrt(2), dtype)
+    sqrt2 = np.array(np.sqrt(2), dtype=dtype)
     l = self.__norm_cdf((lower - loc) / scale, sqrt2, dtype)
     u = self.__norm_cdf((upper - loc) / scale, sqrt2, dtype)
 
     # Uniformly fill tensor with values from [l, u], then translate to
     # [2l-1, 2u-1].
     key = self.split_key() if key is None else _formalize_key(key)
-    out = jr.uniform(key, size, dtype,
-                     minval=lax.nextafter(2 * l - 1, np.array(np.inf, dtype=dtype)),
-                     maxval=lax.nextafter(2 * u - 1, np.array(-np.inf, dtype=dtype)))
+    out = uniform_for_unit(
+      key, size, dtype,
+      minval=lax.nextafter(2 * l - 1, np.array(np.inf, dtype=dtype)),
+      maxval=lax.nextafter(2 * u - 1, np.array(-np.inf, dtype=dtype))
+    )
 
     # Use inverse cdf transform for normal distribution to get truncated
     # standard normal
@@ -523,7 +538,7 @@ class RandomState(State):
     out = jnp.clip(out,
                    lax.nextafter(lax.stop_gradient(lower), np.array(np.inf, dtype=dtype)),
                    lax.nextafter(lax.stop_gradient(upper), np.array(-np.inf, dtype=dtype)))
-    return out
+    return out if dim == bu.DIMENSIONLESS else bu.Quantity(out, dim=dim)
 
   def _check_p(self, p):
     raise ValueError(f'Parameter p should be within [0, 1], but we got {p}')
@@ -548,6 +563,13 @@ class RandomState(State):
                 dtype: DTypeLike = None):
     mean = _check_py_seq(mean)
     sigma = _check_py_seq(sigma)
+    mean = bu.math.asarray(mean, dtype=dtype)
+    sigma = bu.math.asarray(sigma, dtype=dtype)
+    bu.fail_for_dimension_mismatch(mean, sigma)
+    dim = mean.dim if isinstance(mean, bu.Quantity) else bu.DIMENSIONLESS
+    mean = mean.value if isinstance(mean, bu.Quantity) else mean
+    sigma = sigma.value if isinstance(sigma, bu.Quantity) else sigma
+
     if size is None:
       size = jnp.broadcast_shapes(jnp.shape(mean),
                                   jnp.shape(sigma))
@@ -556,7 +578,7 @@ class RandomState(State):
     samples = jr.normal(key, shape=_size2shape(size), dtype=dtype)
     samples = _loc_scale(mean, sigma, samples)
     samples = jnp.exp(samples)
-    return samples
+    return samples if dim == bu.DIMENSIONLESS else bu.Quantity(samples, dim=dim)
 
   def binomial(self,
                n,
@@ -614,7 +636,7 @@ class RandomState(State):
       size = jnp.shape(p)
     key = self.split_key() if key is None else _formalize_key(key)
     dtype = dtype or environ.dftype()
-    u = jr.uniform(key, size, dtype=dtype)
+    u = uniform_for_unit(key, size, dtype=dtype)
     r = jnp.floor(jnp.log1p(-u) / jnp.log1p(-p))
     return r
 
@@ -640,20 +662,28 @@ class RandomState(State):
     dtype = dtype or environ.ditype()
     return jnp.asarray(r, dtype=dtype)
 
-  def multivariate_normal(self,
-                          mean,
-                          cov,
-                          size: Optional[Size] = None,
-                          method: str = 'cholesky',
-                          key: Optional[SeedOrKey] = None,
-                          dtype: DTypeLike = None):
+  def multivariate_normal(
+      self,
+      mean,
+      cov,
+      size: Optional[Size] = None,
+      method: str = 'cholesky',
+      key: Optional[SeedOrKey] = None,
+      dtype: DTypeLike = None
+  ):
     if method not in {'svd', 'eigh', 'cholesky'}:
       raise ValueError("method must be one of {'svd', 'eigh', 'cholesky'}")
     dtype = dtype or environ.dftype()
-    mean = jnp.asarray(_check_py_seq(mean), dtype=dtype)
-    cov = jnp.asarray(_check_py_seq(cov), dtype=dtype)
-    key = self.split_key() if key is None else _formalize_key(key)
+    mean = bu.math.asarray(_check_py_seq(mean), dtype=dtype)
+    cov = bu.math.asarray(_check_py_seq(cov), dtype=dtype)
+    if isinstance(mean, bu.Quantity):
+      assert isinstance(cov, bu.Quantity)
+      assert mean.dim ** 2 == cov.dim
+    mean = mean.value if isinstance(mean, bu.Quantity) else mean
+    cov = cov.value if isinstance(cov, bu.Quantity) else cov
+    dim = mean.dim if isinstance(mean, bu.Quantity) else bu.DIMENSIONLESS
 
+    key = self.split_key() if key is None else _formalize_key(key)
     if not jnp.ndim(mean) >= 1:
       raise ValueError(f"multivariate_normal requires mean.ndim >= 1, got mean.ndim == {jnp.ndim(mean)}")
     if not jnp.ndim(cov) >= 2:
@@ -678,7 +708,7 @@ class RandomState(State):
       factor = jnp.linalg.cholesky(cov)
     normal_samples = jr.normal(key, size + mean.shape[-1:], dtype=dtype)
     r = mean + jnp.einsum('...ij,...j->...i', factor, normal_samples)
-    return r
+    return r if dim == bu.DIMENSIONLESS else bu.Quantity(r, dim=dim)
 
   def rayleigh(self,
                scale=1.0,
@@ -690,7 +720,7 @@ class RandomState(State):
       size = jnp.shape(scale)
     key = self.split_key() if key is None else _formalize_key(key)
     dtype = dtype or environ.dftype()
-    x = jnp.sqrt(-2. * jnp.log(jr.uniform(key, shape=_size2shape(size), minval=0, maxval=1, dtype=dtype)))
+    x = jnp.sqrt(-2. * jnp.log(uniform_for_unit(key, shape=_size2shape(size), minval=0, maxval=1, dtype=dtype)))
     r = x * scale
     return r
 
@@ -734,7 +764,7 @@ class RandomState(State):
         raise ValueError(f'"a" should be a scalar when "size" is provided. But we got {a}')
     size = _size2shape(size)
     dtype = dtype or environ.dftype()
-    random_uniform = jr.uniform(key=key, shape=size, minval=0, maxval=1, dtype=dtype)
+    random_uniform = uniform_for_unit(key=key, shape=size, minval=0, maxval=1, dtype=dtype)
     r = jnp.power(-jnp.log1p(-random_uniform), 1.0 / a)
     return r
 
@@ -754,7 +784,7 @@ class RandomState(State):
         raise ValueError(f'"a" should be a scalar when "size" is provided. But we got {a}')
     size = _size2shape(size)
     dtype = dtype or environ.dftype()
-    random_uniform = jr.uniform(key=key, shape=size, minval=0, maxval=1, dtype=dtype)
+    random_uniform = uniform_for_unit(key=key, shape=size, minval=0, maxval=1, dtype=dtype)
     r = jnp.power(-jnp.log1p(-random_uniform), 1.0 / a)
     if scale is not None:
       r /= scale

brainstate/transform/_jit_error.py

@@ -15,14 +15,14 @@
 
 from __future__ import annotations
 
-from functools import wraps, partial
+import functools
+from functools import partial
 from typing import Callable, Union
 
 import jax
 from jax import numpy as jnp
 from jax.core import Primitive, ShapedArray
-from jax.interpreters import batching, mlir, xla
-from jax.lax import cond
+from jax.interpreters import batching, mlir
 
 from brainstate._utils import set_module_as
 
@@ -32,16 +32,39 @@ __all__ = [
 
 
 @set_module_as('brainstate.transform')
-def remove_vmap(x, op='any'):
+def remove_vmap(x, op: str = 'any'):
   if op == 'any':
     return _any_without_vmap(x)
   elif op == 'all':
     return _all_without_vmap(x)
+  elif op == 'none':
+    return _without_vmap(x)
   else:
     raise ValueError(f'Do not support type: {op}')
 
 
-_any_no_vmap_prim = Primitive('any_no_vmap')
+def _without_vmap(x):
+  return _no_vmap_prim.bind(x)
+
+
+def _without_vmap_imp(x):
+  return x
+
+
+def _without_vmap_abs(x):
+  return x
+
+
+def _without_vmap_batch(x, batch_axes):
+  (x,) = x
+  return _without_vmap(x), batching.not_mapped
+
+
+_no_vmap_prim = Primitive('no_vmap')
+_no_vmap_prim.def_impl(_without_vmap_imp)
+_no_vmap_prim.def_abstract_eval(_without_vmap_abs)
+batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch
+mlir.register_lowering(_no_vmap_prim, mlir.lower_fun(_without_vmap_imp, multiple_results=False))
 
 
 def _any_without_vmap(x):
@@ -61,16 +84,12 @@ def _any_without_vmap_batch(x, batch_axes):
   return _any_without_vmap(x), batching.not_mapped
 
 
+_any_no_vmap_prim = Primitive('any_no_vmap')
 _any_no_vmap_prim.def_impl(_any_without_vmap_imp)
 _any_no_vmap_prim.def_abstract_eval(_any_without_vmap_abs)
 batching.primitive_batchers[_any_no_vmap_prim] = _any_without_vmap_batch
-if hasattr(xla, "lower_fun"):
-  xla.register_translation(_any_no_vmap_prim,
-                           xla.lower_fun(_any_without_vmap_imp, multiple_results=False, new_style=True))
 mlir.register_lowering(_any_no_vmap_prim, mlir.lower_fun(_any_without_vmap_imp, multiple_results=False))
 
-_all_no_vmap_prim = Primitive('all_no_vmap')
-
 
 def _all_without_vmap(x):
   return _all_no_vmap_prim.bind(x)
@@ -89,50 +108,54 @@ def _all_without_vmap_batch(x, batch_axes):
   return _all_without_vmap(x), batching.not_mapped
 
 
+_all_no_vmap_prim = Primitive('all_no_vmap')
 _all_no_vmap_prim.def_impl(_all_without_vmap_imp)
 _all_no_vmap_prim.def_abstract_eval(_all_without_vmap_abs)
 batching.primitive_batchers[_all_no_vmap_prim] = _all_without_vmap_batch
-if hasattr(xla, "lower_fun"):
-  xla.register_translation(_all_no_vmap_prim,
-                           xla.lower_fun(_all_without_vmap_imp, multiple_results=False, new_style=True))
 mlir.register_lowering(_all_no_vmap_prim, mlir.lower_fun(_all_without_vmap_imp, multiple_results=False))
 
 
-def _err_jit_true_branch(err_fun, x):
-  jax.debug.callback(err_fun, x)
-  return
+def _err_jit_true_branch(err_fun, args, kwargs):
+  jax.debug.callback(err_fun, *args, **kwargs)
 
 
-def _err_jit_false_branch(x):
-  return
+def _err_jit_false_branch(args, kwargs):
+  pass
 
 
-def _cond(err_fun, pred, err_arg):
-  @wraps(err_fun)
-  def true_err_fun(*arg):
-    err_fun(*arg)
+def _error_msg(msg, *arg, **kwargs):
+  if len(arg):
+    msg = msg % arg
+  if len(kwargs):
+    msg = msg.format(**kwargs)
+  raise ValueError(msg)
 
-  cond(pred,
-       partial(_err_jit_true_branch, true_err_fun),
-       _err_jit_false_branch,
-       err_arg)
 
+@set_module_as('brainstate.transform')
+def jit_error(
+    pred,
+    err_fun: Union[Callable, str],
+    *err_args,
+    **err_kwargs,
+):
+  """
+  Check errors in a jit function.
 
-def _error_msg(msg, *arg):
-  if len(arg) == 0:
-    raise ValueError(msg)
-  else:
-    raise ValueError(msg.format(arg))
+  Examples
+  --------
 
+  It can give a function which receive arguments that passed from the JIT variables and raise errors.
 
-@set_module_as('brainstate.transform')
-def jit_error(pred, err_fun: Union[Callable, str], err_arg=None, scope: str = 'any'):
-  """Check errors in a jit function.
+  >>> def error(x):
+  >>>    raise ValueError(f'error {x}')
+  >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,))
+  >>> jit_error(x.sum() < 5., error, x)
+
+  Or, it can be a simple string message.
 
-  >>> def error(arg):
-  >>>    raise ValueError(f'error {arg}')
   >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,))
-  >>> jit_error(x.sum() < 5., error, err_arg=x)
+  >>> jit_error(x.sum() < 5., "Error: the sum is less than 5. Got {s}", s=x.sum())
+
 
   Parameters
   ----------
@@ -140,19 +163,18 @@ def jit_error(pred, err_fun: Union[Callable, str], err_arg=None, scope: str = 'a
     The boolean prediction.
   err_fun: callable
     The error function, which raise errors.
-  err_arg: any
+  err_args: 
     The arguments which passed into `err_f`.
-  scope: str
-    The scope of the error message. Can be None, 'all' or 'any'.
+  err_kwargs: 
+    The keywords which passed into `err_f`.
   """
   if isinstance(err_fun, str):
     err_fun = partial(_error_msg, err_fun)
-  if scope is None:
-    pred = pred
-  elif scope == 'all':
-    pred = remove_vmap(pred, 'all')
-  elif scope == 'any':
-    pred = remove_vmap(pred, 'any')
-  else:
-    raise ValueError(f"Unknown scope: {scope}")
-  _cond(err_fun, pred, err_arg)
+
+  jax.lax.cond(
+    remove_vmap(pred, op='any'),
+    partial(_err_jit_true_branch, err_fun),
+    _err_jit_false_branch,
+    jax.tree.map(functools.partial(remove_vmap, op='none'), err_args),
+    jax.tree.map(functools.partial(remove_vmap, op='none'), err_kwargs),
+  )

brainstate/transform/_jit_error_test.py

@@ -0,0 +1,55 @@
+# 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
+import jaxlib.xla_extension
+import jax.numpy as jnp
+
+import brainstate as bst
+
+
+class TestJitError(unittest.TestCase):
+  def test1(self):
+    with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+      bst.transform.jit_error(True, 'error')
+
+    def err_f(x):
+      raise ValueError(f'error: {x}')
+
+    with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+      bst.transform.jit_error(True, err_f, 1.)
+
+  def test_vmap(self):
+
+    def f(x):
+      bst.transform.jit_error(x, 'error: {x}', x=x)
+
+    jax.vmap(f)(jnp.array([False, False, False]))
+    with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+      jax.vmap(f)(jnp.array([True, False, False]))
+
+  def test_vmap_vmap(self):
+
+    def f(x):
+      bst.transform.jit_error(x, 'error: {x}', x=x)
+
+    jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
+                                     [False, False, False]]))
+    with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+      jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
+                                       [True, False, False]]))
+

brainstate/transform/_make_jaxpr.py

@@ -369,11 +369,11 @@ class StatefulFunction(object):
     self._state_trace[cache_key] = _state_trace
     with _state_trace:
       out = self.fun(*args, **kwargs)
-      state_values = _state_trace.collect_values('read', 'write')
+      state_values = _state_trace.collect_values('read', 'write', check_val_tree=True)
     _state_trace.recovery_original_values()
 
-    # return states is not allowed
-    # checking whether the states are returned
+    # State instance as functional returns 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}"))
@@ -399,7 +399,6 @@ 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) = _make_jaxpr(
           functools.partial(self._wrapped_fun_to_eval, cache_key),
           static_argnums=self.static_argnums,
@@ -435,8 +434,11 @@ class StatefulFunction(object):
     """
     # state checking
     cache_key = self.get_arg_cache_key(*args, **kwargs)
-    states = self.get_states(cache_key)
+    states: Sequence[State] = self.get_states(cache_key)
     assert len(state_vals) == len(states), 'State length mismatch.'
+    # # No need to check, because the make_jaxpr() has been checked whether the value's tree is correct.
+    # for val, st in zip(state_vals, states):  # check state's value tree structure
+    #   st._check_value_tree(val)
 
     # parameters
     args = tuple(args[i] for i in range(len(args)) if i not in self.static_argnums)
@@ -450,6 +452,9 @@ class StatefulFunction(object):
     # output processing
     out, new_state_vals = out_treedef.unflatten(jaxpr_outs)
     assert len(new_state_vals) == len(state_vals), 'State length mismatch.'
+    # # No need to check, because the make_jaxpr() has been checked whether the value's tree is correct.
+    # for val, st in zip(new_state_vals, states):  # check state's value tree structure
+    #   st._check_value_tree(val)
     return new_state_vals, out
 
   def jaxpr_call_auto(self, *args, **kwargs) -> Any:

brainstate/typing.py

@@ -16,6 +16,7 @@
 
 from typing import Any, Sequence, Protocol, Union
 
+import brainunit as bu
 import jax
 import numpy as np
 
@@ -43,6 +44,7 @@ ArrayLike = Union[
   np.ndarray,  # NumPy array type
   np.bool_, np.number,  # NumPy scalar types
   bool, int, float, complex,  # Python scalar types
+  bu.Quantity,  # quantity
 ]
 
 # --- Dtype --- #

{brainstate-0.0.1.post20240623.dist-info → brainstate-0.0.1.1.post20240708.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.1.post20240623
+Version: 0.0.1.1.post20240708
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BrainPy Team