brainstate 0.0.2.post20241009__py2.py3-none-any.whl → 0.1.0__py2.py3-none-any.whl

brainstate/transform/_conditions.py

@@ -1,334 +0,0 @@
-# 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.
-# ==============================================================================
-
-from __future__ import annotations
-
-import operator
-from collections.abc import Callable, Sequence
-from functools import wraps, reduce
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from brainstate._utils import set_module_as
-from ._error_if import jit_error_if
-from ._make_jaxpr import StatefulFunction, _assign_state_values
-
-__all__ = [
-  'cond', 'switch', 'ifelse',
-]
-
-
-def _wrapped_fun(stateful_fun: StatefulFunction, states, return_states=True):
-  @wraps(stateful_fun.fun)
-  def wrapped_branch(state_vals, *operands):
-    assert len(states) == len(state_vals)
-    for st, val in zip(states, state_vals):
-      st.value = val
-    out = stateful_fun.jaxpr_call_auto(*operands)
-    if return_states:
-      return tuple(st.value for st in states), out
-    return out
-
-  return wrapped_branch
-
-
-@set_module_as('brainstate.transform')
-def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
-  """
-  Conditionally apply ``true_fun`` or ``false_fun``.
-
-  Provided arguments are correctly typed, ``cond()`` has equivalent
-  semantics to this Python implementation, where ``pred`` must be a
-  scalar type::
-
-    def cond(pred, true_fun, false_fun, *operands):
-      if pred:
-        return true_fun(*operands)
-      else:
-        return false_fun(*operands)
-
-
-  In contrast with :func:`jax.lax.select`, using ``cond`` indicates that only one of
-  the two branches is executed (up to compiler rewrites and optimizations).
-  However, when transformed with :func:`~jax.vmap` to operate over a batch of
-  predicates, ``cond`` is converted to :func:`~jax.lax.select`.
-
-  Args:
-    pred: Boolean scalar type, indicating which branch function to apply.
-    true_fun: Function (A -> B), to be applied if ``pred`` is True.
-    false_fun: Function (A -> B), to be applied if ``pred`` is False.
-    operands: Operands (A) input to either branch depending on ``pred``. The
-      type can be a scalar, array, or any pytree (nested Python tuple/list/dict)
-      thereof.
-
-  Returns:
-    Value (B) of either ``true_fun(*operands)`` or ``false_fun(*operands)``,
-    depending on the value of ``pred``. The type can be a scalar, array, or any
-    pytree (nested Python tuple/list/dict) thereof.
-  """
-  if not (callable(true_fun) and callable(false_fun)):
-    raise TypeError("true_fun and false_fun arguments should be callable.")
-
-  if pred is None:
-    raise TypeError("cond predicate is None")
-  if isinstance(pred, Sequence) or np.ndim(pred) != 0:
-    raise TypeError(f"Pred must be a scalar, got {pred} of " +
-                    (f"type {type(pred)}" if isinstance(pred, Sequence)
-                     else f"shape {np.shape(pred)}."))
-
-  # check pred
-  try:
-    pred_dtype = jax.dtypes.result_type(pred)
-  except TypeError as err:
-    raise TypeError("Pred type must be either boolean or number, got {}.".format(pred)) from err
-  if pred_dtype.kind != 'b':
-    if pred_dtype.kind in 'iuf':
-      pred = pred != 0
-    else:
-      raise TypeError("Pred type must be either boolean or number, got {}.".format(pred_dtype))
-
-  # not jit
-  if jax.config.jax_disable_jit and isinstance(jax.core.get_aval(pred), jax.core.ConcreteArray):
-    if pred:
-      return true_fun(*operands)
-    else:
-      return false_fun(*operands)
-
-  # evaluate jaxpr
-  true_fun_wrap = StatefulFunction(true_fun).make_jaxpr(*operands)
-  false_fun_wrap = StatefulFunction(false_fun).make_jaxpr(*operands)
-
-  # wrap the functions
-  all_states = tuple(set(true_fun_wrap.get_states() + false_fun_wrap.get_states()))
-  true_fun = _wrapped_fun(true_fun_wrap, all_states)
-  false_fun = _wrapped_fun(false_fun_wrap, all_states)
-
-  # operands
-  operands = ([st.value for st in all_states],) + operands
-
-  # cond
-  state_vals, out = jax.lax.cond(pred, true_fun, false_fun, *operands)
-  _assign_state_values(all_states, state_vals)
-  return out
-
-  # ops, ops_tree = jax.tree.flatten(operands)
-  # linear_ops = [False] * len(ops)
-  # ops_avals = tuple(jax.util.safe_map(_abstractify, ops))
-  #
-  # # true and false jaxprs
-  # jaxprs, consts, out_trees = _initial_style_jaxprs_with_common_consts(
-  #   (true_fun, false_fun), ops_tree, ops_avals, 'cond')
-  # if any(isinstance(op_aval, state.AbstractRef) for op_aval in ops_avals):
-  #   raise ValueError("Cannot pass `Ref`s into `cond`.")
-  # true_jaxpr, false_jaxpr = jaxprs
-  # out_tree, false_out_tree = out_trees
-  # if any(isinstance(out_aval, state.AbstractRef) for out_aval in true_jaxpr.out_avals + false_jaxpr.out_avals):
-  #   raise ValueError("Cannot return `Ref`s from `cond`.")
-  #
-  # _check_tree_and_avals("true_fun and false_fun output",
-  #                       out_tree, true_jaxpr.out_avals,
-  #                       false_out_tree, false_jaxpr.out_avals)
-  # joined_effects = jax.core.join_effects(true_jaxpr.effects, false_jaxpr.effects)
-  # disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(joined_effects)
-  # if disallowed_effects:
-  #   raise NotImplementedError(f'Effects not supported in `cond`: {disallowed_effects}')
-  #
-  # # replace jaxpr effects
-  # index = jax.lax.convert_element_type(pred, np.int32)
-  # if joined_effects:
-  #   # Raise index in case of effects to allow data-dependence-based discharging
-  #   # of those effects (even if they don't have an explicit data dependence).
-  #   index = jax.core.raise_as_much_as_possible(index)
-  # false_jaxpr = _replace_jaxpr_effects(false_jaxpr, joined_effects)
-  # true_jaxpr = _replace_jaxpr_effects(true_jaxpr, joined_effects)
-  #
-  # # bind
-  # linear = [False] * len(consts) + linear_ops
-  # cond_outs = jax.lax.cond_p.bind(index, *consts, *ops, branches=(false_jaxpr, true_jaxpr), linear=tuple(linear))
-  #
-  # # outputs
-  # st_vals, out = jax.tree.unflatten(out_tree, cond_outs)
-  # for st, val in zip(all_states, st_vals):
-  #   st.value = val
-  # return out
-
-
-@set_module_as('brainstate.transform')
-def switch(index, branches: Sequence[Callable], *operands):
-  """
-  Apply exactly one of ``branches`` given by ``index``.
-
-  If ``index`` is out of bounds, it is clamped to within bounds.
-
-  Has the semantics of the following Python::
-
-    def switch(index, branches, *operands):
-      index = clamp(0, index, len(branches) - 1)
-      return branches[index](*operands)
-
-  Internally this wraps XLA's `Conditional
-  <https://www.tensorflow.org/xla/operation_semantics#conditional>`_
-  operator. However, when transformed with :func:`~jax.vmap` to operate over a
-  batch of predicates, ``cond`` is converted to :func:`~jax.lax.select`.
-
-  Args:
-    index: Integer scalar type, indicating which branch function to apply.
-    branches: Sequence of functions (A -> B) to be applied based on ``index``.
-    operands: Operands (A) input to whichever branch is applied.
-
-  Returns:
-    Value (B) of ``branch(*operands)`` for the branch that was selected based
-    on ``index``.
-  """
-  # check branches
-  if not all(callable(branch) for branch in branches):
-    raise TypeError("branches argument should be a sequence of callables.")
-
-  # check index
-  if len(np.shape(index)) != 0:
-    raise TypeError(f"Branch index must be scalar, got {index} of shape {np.shape(index)}.")
-  try:
-    index_dtype = jax.dtypes.result_type(index)
-  except TypeError as err:
-    msg = f"Index type must be an integer, got {index}."
-    raise TypeError(msg) from err
-  if index_dtype.kind not in 'iu':
-    raise TypeError(f"Index type must be an integer, got {index} as {index_dtype}")
-
-  # format branches
-  branches = tuple(branches)
-  if len(branches) == 0:
-    raise ValueError("Empty branch sequence")
-  elif len(branches) == 1:
-    return branches[0](*operands)
-
-  # format index
-  index = jax.lax.convert_element_type(index, np.int32)
-  lo = np.array(0, np.int32)
-  hi = np.array(len(branches) - 1, np.int32)
-  index = jax.lax.clamp(lo, index, hi)
-
-  # not jit
-  if jax.config.jax_disable_jit and isinstance(jax.core.core.get_aval(index), jax.core.ConcreteArray):
-    return branches[int(index)](*operands)
-
-  # evaluate jaxpr
-  wrapped_branches = [StatefulFunction(branch) for branch in branches]
-  for wrapped_branch in wrapped_branches:
-    wrapped_branch.make_jaxpr(*operands)
-
-  # wrap the functions
-  all_states = tuple(set(reduce(operator.add, [wrapped_branch.get_states() for wrapped_branch in wrapped_branches])))
-  branches = tuple(_wrapped_fun(wrapped_branch, all_states) for wrapped_branch in wrapped_branches)
-
-  # operands
-  operands = ([st.value for st in all_states],) + operands
-
-  # switch
-  state_vals, out = jax.lax.switch(index, branches, *operands)
-  _assign_state_values(all_states, state_vals)
-  return out
-
-  # ops, ops_tree = jax.tree.flatten(operands)
-  # ops_avals = tuple(jax.util.safe_map(_abstractify, ops))
-  #
-  # # true jaxprs
-  # jaxprs, consts, out_trees = _initial_style_jaxprs_with_common_consts(
-  #   branches, ops_tree, ops_avals, primitive_name='switch')
-  # for i, (out_tree, jaxpr) in enumerate(zip(out_trees[1:], jaxprs[1:])):
-  #   _check_tree_and_avals(f"branch 0 and {i + 1} outputs",
-  #                         out_trees[0], jaxprs[0].out_avals,
-  #                         out_tree, jaxpr.out_avals)
-  # joined_effects = jax.core.join_effects(*(jaxpr.effects for jaxpr in jaxprs))
-  # disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(joined_effects)
-  # if disallowed_effects:
-  #   raise NotImplementedError(f'Effects not supported in `switch`: {disallowed_effects}')
-  # if joined_effects:
-  #   # Raise index in case of effects to allow data-dependence-based discharging
-  #   # of those effects (even if they don't have an explicit data dependence).
-  #   index = jax.core.raise_as_much_as_possible(index)
-  #
-  # # bind
-  # linear = (False,) * (len(consts) + len(ops))
-  # cond_outs = jax.lax.cond_p.bind(index, *consts, *ops, branches=tuple(jaxprs), linear=linear)
-  #
-  # # outputs
-  # st_vals, out = jax.tree.unflatten(out_trees[0], cond_outs)
-  # for st, val in zip(all_states, st_vals):
-  #   st.value = val
-  # return out
-
-
-@set_module_as('brainstate.transform')
-def ifelse(conditions, branches, *operands, check_cond: bool = True):
-  """
-  ``If-else`` control flows looks like native Pythonic programming.
-
-  Examples
-  --------
-
-  >>> import brainstate as bst
-  >>> def f(a):
-  >>>    return bst.transform.ifelse(conditions=[a > 10, a > 5, a > 2, a > 0],
-  >>>                               branches=[lambda: 1,
-  >>>                                         lambda: 2,
-  >>>                                         lambda: 3,
-  >>>                                         lambda: 4,
-  >>>                                         lambda: 5])
-  >>> f(1)
-  4
-  >>> f(0)
-  5
-
-  Parameters
-  ----------
-  conditions: bool, sequence of bool, Array
-    The boolean conditions.
-  branches: Any
-    The branches, at least has two elements. Elements can be functions,
-    arrays, or numbers. The number of ``branches`` and ``conditions`` has
-    the relationship of `len(branches) == len(conditions) + 1`.
-    Each branch should receive one arguement for ``operands``.
-  *operands: optional, Any
-    The operands for each branch.
-  check_cond: bool
-    Whether to check the conditions. Default is True.
-
-  Returns
-  -------
-  res: Any
-    The results of the control flow.
-  """
-  # check branches
-  if not all(callable(branch) for branch in branches):
-    raise TypeError("branches argument should be a sequence of callables.")
-
-  # format branches
-  branches = tuple(branches)
-  if len(branches) == 0:
-    raise ValueError("Empty branch sequence")
-  elif len(branches) == 1:
-    return branches[0](*operands)
-  if len(conditions) != len(branches):
-    raise ValueError("The number of conditions should be equal to the number of branches.")
-
-  # format index
-  conditions = jnp.asarray(conditions, np.int32)
-  if check_cond:
-    jit_error_if(jnp.sum(conditions) != 1, "Only one condition can be True. But got {}.", err_arg=conditions)
-  index = jnp.where(conditions, size=1, fill_value=len(conditions) - 1)[0][0]
-  return switch(index, branches, *operands)

brainstate/transform/_conditions_test.py

@@ -1,220 +0,0 @@
-# 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 jax.numpy as jnp
-
-import brainstate as bst
-
-
-class TestCond(unittest.TestCase):
-  def test1(self):
-    bst.random.seed(1)
-    bst.transform.cond(True, lambda: bst.random.random(10), lambda: bst.random.random(10))
-    bst.transform.cond(False, lambda: bst.random.random(10), lambda: bst.random.random(10))
-
-  def test2(self):
-    st1 = bst.State(bst.random.rand(10))
-    st2 = bst.State(bst.random.rand(2))
-    st3 = bst.State(bst.random.rand(5))
-    st4 = bst.State(bst.random.rand(2, 10))
-
-    def true_fun(x):
-      st1.value = st2.value @ st4.value + x
-
-    def false_fun(x):
-      st3.value = (st3.value + 1.) * x
-
-    bst.transform.cond(True, true_fun, false_fun, 2.)
-    assert not isinstance(st1.value, jax.core.Tracer)
-    assert not isinstance(st2.value, jax.core.Tracer)
-    assert not isinstance(st3.value, jax.core.Tracer)
-    assert not isinstance(st4.value, jax.core.Tracer)
-
-
-class TestSwitch(unittest.TestCase):
-  def testSwitch(self):
-    def branch(x):
-      y = jax.lax.mul(2, x)
-      return y, jax.lax.mul(2, y)
-
-    branches = [lambda x: (x, x),
-                branch,
-                lambda x: (x, -x)]
-
-    def fun(x):
-      if x <= 0:
-        return branches[0](x)
-      elif x == 1:
-        return branches[1](x)
-      else:
-        return branches[2](x)
-
-    def cfun(x):
-      return bst.transform.switch(x, branches, x)
-
-    self.assertEqual(fun(-1), cfun(-1))
-    self.assertEqual(fun(0), cfun(0))
-    self.assertEqual(fun(1), cfun(1))
-    self.assertEqual(fun(2), cfun(2))
-    self.assertEqual(fun(3), cfun(3))
-
-    cfun = jax.jit(cfun)
-
-    self.assertEqual(fun(-1), cfun(-1))
-    self.assertEqual(fun(0), cfun(0))
-    self.assertEqual(fun(1), cfun(1))
-    self.assertEqual(fun(2), cfun(2))
-    self.assertEqual(fun(3), cfun(3))
-
-  def testSwitchMultiOperands(self):
-    branches = [jax.lax.add, jax.lax.mul]
-
-    def fun(x):
-      i = 0 if x <= 0 else 1
-      return branches[i](x, x)
-
-    def cfun(x):
-      return bst.transform.switch(x, branches, x, x)
-
-    self.assertEqual(fun(-1), cfun(-1))
-    self.assertEqual(fun(0), cfun(0))
-    self.assertEqual(fun(1), cfun(1))
-    self.assertEqual(fun(2), cfun(2))
-    cfun = jax.jit(cfun)
-    self.assertEqual(fun(-1), cfun(-1))
-    self.assertEqual(fun(0), cfun(0))
-    self.assertEqual(fun(1), cfun(1))
-    self.assertEqual(fun(2), cfun(2))
-
-  def testSwitchResidualsMerge(self):
-    def get_conds(fun):
-      jaxpr = jax.make_jaxpr(jax.grad(fun))(0., 0)
-      return [eqn for eqn in jaxpr.jaxpr.eqns if eqn.primitive.name == 'cond']
-
-    def branch_invars_len(cond_eqn):
-      lens = [len(jaxpr.jaxpr.invars) for jaxpr in cond_eqn.params['branches']]
-      assert len(set(lens)) == 1
-      return lens[0]
-
-    def branch_outvars_len(cond_eqn):
-      lens = [len(jaxpr.jaxpr.outvars) for jaxpr in cond_eqn.params['branches']]
-      assert len(set(lens)) == 1
-      return lens[0]
-
-    branches1 = [lambda x: jnp.sin(x),
-                 lambda x: jnp.cos(x)]  # branch residuals overlap, should be reused
-    branches2 = branches1 + [lambda x: jnp.sinh(x)]  # another overlapping residual, expect reuse
-    branches3 = branches2 + [lambda x: jnp.sin(x) + jnp.cos(x)]  # requires one more residual slot
-
-    def fun1(x, i):
-      return bst.transform.switch(i + 1, branches1, x)
-
-    def fun2(x, i):
-      return bst.transform.switch(i + 1, branches2, x)
-
-    def fun3(x, i):
-      return bst.transform.switch(i + 1, branches3, x)
-
-    fwd1, bwd1 = get_conds(fun1)
-    fwd2, bwd2 = get_conds(fun2)
-    fwd3, bwd3 = get_conds(fun3)
-
-    fwd1_num_out = branch_outvars_len(fwd1)
-    fwd2_num_out = branch_outvars_len(fwd2)
-    fwd3_num_out = branch_outvars_len(fwd3)
-    assert fwd1_num_out == fwd2_num_out
-    assert fwd3_num_out == fwd2_num_out + 1
-
-    bwd1_num_in = branch_invars_len(bwd1)
-    bwd2_num_in = branch_invars_len(bwd2)
-    bwd3_num_in = branch_invars_len(bwd3)
-    assert bwd1_num_in == bwd2_num_in
-    assert bwd3_num_in == bwd2_num_in + 1
-
-  def testOneBranchSwitch(self):
-    branch = lambda x: -x
-    f = lambda i, x: bst.transform.switch(i, [branch], x)
-    x = 7.
-    self.assertEqual(f(-1, x), branch(x))
-    self.assertEqual(f(0, x), branch(x))
-    self.assertEqual(f(1, x), branch(x))
-    cf = jax.jit(f)
-    self.assertEqual(cf(-1, x), branch(x))
-    self.assertEqual(cf(0, x), branch(x))
-    self.assertEqual(cf(1, x), branch(x))
-    cf = jax.jit(f, static_argnums=0)
-    self.assertEqual(cf(-1, x), branch(x))
-    self.assertEqual(cf(0, x), branch(x))
-    self.assertEqual(cf(1, x), branch(x))
-
-
-class TestIfElse(unittest.TestCase):
-  def test1(self):
-    def f(a):
-      return bst.transform.ifelse(conditions=[a < 0,
-                                              a >= 0 and a < 2,
-                                              a >= 2 and a < 5,
-                                              a >= 5 and a < 10,
-                                              a >= 10],
-                                  branches=[lambda: 1,
-                                           lambda: 2,
-                                           lambda: 3,
-                                           lambda: 4,
-                                           lambda: 5])
-
-    self.assertTrue(f(3) == 3)
-    self.assertTrue(f(1) == 2)
-    self.assertTrue(f(-1) == 1)
-
-  def test_vmap(self):
-    def f(operands):
-      f = lambda a: bst.transform.ifelse([a > 10,
-                                          jnp.logical_and(a <= 10, a > 5),
-                                          jnp.logical_and(a <= 5, a > 2),
-                                          jnp.logical_and(a <= 2, a > 0),
-                                          a <= 0],
-                                         [lambda _: 1,
-                                         lambda _: 2,
-                                         lambda _: 3,
-                                         lambda _: 4,
-                                         lambda _: 5, ],
-                                         a)
-      return jax.vmap(f)(operands)
-
-    r = f(bst.random.randint(-20, 20, 200))
-    self.assertTrue(r.size == 200)
-
-  def test_grad1(self):
-    def F2(x):
-      return bst.transform.ifelse((x >= 10, x < 10),
-                                  [lambda x: x, lambda x: x ** 2, ],
-                                  x)
-
-    self.assertTrue(jax.grad(F2)(9.0) == 18.)
-    self.assertTrue(jax.grad(F2)(11.0) == 1.)
-
-  def test_grad2(self):
-    def F3(x):
-      return bst.transform.ifelse((x >= 10, jnp.logical_and(x >= 0, x < 10), x < 0),
-                                  [lambda x: x,
-                                  lambda x: x ** 2,
-                                  lambda x: x ** 4, ],
-                                  x)
-
-    self.assertTrue(jax.grad(F3)(9.0) == 18.)
-    self.assertTrue(jax.grad(F3)(11.0) == 1.)

brainstate/transform/_error_if.py

@@ -1,94 +0,0 @@
-# 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.
-# ==============================================================================
-
-from __future__ import annotations
-
-import functools
-from functools import partial
-from typing import Callable, Union
-
-import jax
-
-from brainstate._utils import set_module_as
-from ._unvmap import unvmap
-
-__all__ = [
-  'jit_error_if',
-]
-
-
-def _err_jit_true_branch(err_fun, args, kwargs):
-  jax.debug.callback(err_fun, *args, **kwargs)
-
-
-def _err_jit_false_branch(args, kwargs):
-  pass
-
-
-def _error_msg(msg, *arg, **kwargs):
-  if len(arg):
-    msg = msg % arg
-  if len(kwargs):
-    msg = msg.format(**kwargs)
-  raise ValueError(msg)
-
-
-@set_module_as('brainstate.transform')
-def jit_error_if(
-    pred,
-    error: Union[Callable, str],
-    *err_args,
-    **err_kwargs,
-):
-  """
-  Check errors in a jit function.
-
-  Examples
-  --------
-
-  It can give a function which receive arguments that passed from the JIT variables and raise errors.
-
-  >>> def error(x):
-  >>>    raise ValueError(f'error {x}')
-  >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,))
-  >>> jit_error_if(x.sum() < 5., error, x)
-
-  Or, it can be a simple string message.
-
-  >>> x = jax.random.uniform(jax.random.PRNGKey(0), (10,))
-  >>> jit_error_if(x.sum() < 5., "Error: the sum is less than 5. Got {s}", s=x.sum())
-
-
-  Parameters
-  ----------
-  pred: bool, Array
-    The boolean prediction.
-  error: callable, str
-    The error function, which raise errors, or a string indicating the error message.
-  err_args: 
-    The arguments which passed into `err_f`.
-  err_kwargs: 
-    The keywords which passed into `err_f`.
-  """
-  if isinstance(error, str):
-    error = partial(_error_msg, error)
-
-  jax.lax.cond(
-    unvmap(pred, op='any'),
-    partial(_err_jit_true_branch, error),
-    _err_jit_false_branch,
-    jax.tree.map(functools.partial(unvmap, op='none'), err_args),
-    jax.tree.map(functools.partial(unvmap, op='none'), err_kwargs),
-  )

brainstate/transform/_error_if_test.py

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