brainstate 0.0.2.post20240825__py2.py3-none-any.whl → 0.0.2.post20240826__py2.py3-none-any.whl

brainstate/_module.py

@@ -60,7 +60,7 @@ from . import environ
 from ._state import State, StateDictManager, visible_state_dict
 from ._utils import set_module_as
 from .mixin import Mixin, Mode, DelayedInit, JointTypes, Batching, UpdateReturn
-from .transform import jit_error
+from .transform import jit_error_if
 from .typing import Size, ArrayLike, PyTree
 from .util import unique_name, DictManager, get_unique_name
 
@@ -1212,7 +1212,7 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
         raise ValueError(f'The request delay length should be less than the '
                          f'maximum delay {self.max_length - 1}. But we got {delay_len}')
 
-      jit_error(delay_step >= self.max_length, _check_delay, delay_step)
+      jit_error_if(delay_step >= self.max_length, _check_delay, delay_step)
 
     # rotation method
     if self.delay_method == _DELAY_ROTATE:
@@ -1263,10 +1263,10 @@ class Delay(ExtendedUpdateWithBA, DelayedInit):
                          f'[{t_now - self.max_time - dt}, {t_now}], '
                          f'but we got {t_delay}')
 
-      jit_error(jnp.logical_or(delay_time > current_time,
-                               delay_time < current_time - self.max_time - dt),
-                _check_delay,
-                current_time, delay_time)
+      jit_error_if(jnp.logical_or(delay_time > current_time,
+                                  delay_time < current_time - self.max_time - dt),
+                   _check_delay,
+                   current_time, delay_time)
 
     diff = current_time - delay_time
     float_time_step = diff / dt

brainstate/random.py

@@ -33,7 +33,7 @@ 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 .transform._error_if import jit_error_if
 from .typing import DTypeLike, Size, SeedOrKey
 
 __all__ = [
@@ -498,7 +498,7 @@ class RandomState(State):
       bu.Quantity(scale).in_unit(unit).mantissa
     )
 
-    jit_error(
+    jit_error_if(
       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."
@@ -549,7 +549,7 @@ class RandomState(State):
                 size: Optional[Size] = None,
                 key: Optional[SeedOrKey] = None):
     p = _check_py_seq(p)
-    jit_error(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
+    jit_error_if(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
     if size is None:
       size = jnp.shape(p)
     key = self.split_key() if key is None else _formalize_key(key)
@@ -592,7 +592,7 @@ class RandomState(State):
                dtype: DTypeLike = None):
     n = _check_py_seq(n)
     p = _check_py_seq(p)
-    jit_error(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
+    jit_error_if(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
     if size is None:
       size = jnp.broadcast_shapes(jnp.shape(n), jnp.shape(p))
     key = self.split_key() if key is None else _formalize_key(key)
@@ -656,7 +656,7 @@ class RandomState(State):
     key = self.split_key() if key is None else _formalize_key(key)
     n = _check_py_seq(n)
     pvals = _check_py_seq(pvals)
-    jit_error(jnp.sum(pvals[:-1]) > 1., self._check_p2, pvals)
+    jit_error_if(jnp.sum(pvals[:-1]) > 1., self._check_p2, pvals)
     if isinstance(n, jax.core.Tracer):
       raise ValueError("The total count parameter `n` should not be a jax abstract array.")
     size = _size2shape(size)

brainstate/transform/__init__.py

@@ -19,17 +19,27 @@ This module contains the functions for the transformation of the brain data.
 
 from ._autograd import *
 from ._autograd import __all__ as _gradients_all
-from ._control import *
-from ._control import __all__ as _controls_all
+from ._conditions import *
+from ._conditions import __all__ as _conditions_all
+from ._error_if import *
+from ._error_if import __all__ as _jit_error_all
 from ._jit import *
 from ._jit import __all__ as _jit_all
-from ._jit_error import *
-from ._jit_error import __all__ as _jit_error_all
+from ._loop_collect_return import *
+from ._loop_collect_return import __all__ as _loops_all
+from ._loop_no_collection import *
+from ._loop_no_collection import __all__ as _loops_no_collection_all
 from ._make_jaxpr import *
 from ._make_jaxpr import __all__ as _make_jaxpr_all
+from ._mapping import *
+from ._mapping import __all__ as _mapping_all
 from ._progress_bar import *
 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
+__all__ = (_gradients_all + _jit_error_all + _conditions_all + _loops_all +
+           _make_jaxpr_all + _jit_all + _progress_bar_all + _loops_no_collection_all +
+           _mapping_all)
 
-del _gradients_all, _jit_error_all, _controls_all, _make_jaxpr_all, _jit_all, _progress_bar_all
+del (_gradients_all, _jit_error_all, _conditions_all, _loops_all,
+     _make_jaxpr_all, _jit_all, _progress_bar_all, _loops_no_collection_all,
+     _mapping_all)

brainstate/transform/_conditions.py

@@ -0,0 +1,334 @@
+# 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/{_controls_test.py → _conditions_test.py}

@@ -18,20 +18,20 @@ import unittest
 import jax
 import jax.numpy as jnp
 
-import brainstate as bc
+import brainstate as bst
 
 
 class TestCond(unittest.TestCase):
   def test1(self):
-    bc.random.seed(1)
-    bc.transform.cond(True, lambda: bc.random.random(10), lambda: bc.random.random(10))
-    bc.transform.cond(False, lambda: bc.random.random(10), lambda: bc.random.random(10))
+    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 = bc.State(bc.random.rand(10))
-    st2 = bc.State(bc.random.rand(2))
-    st3 = bc.State(bc.random.rand(5))
-    st4 = bc.State(bc.random.rand(2, 10))
+    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
@@ -39,7 +39,7 @@ class TestCond(unittest.TestCase):
     def false_fun(x):
       st3.value = (st3.value + 1.) * x
 
-    bc.transform.cond(True, true_fun, false_fun, 2.)
+    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)
@@ -65,7 +65,7 @@ class TestSwitch(unittest.TestCase):
         return branches[2](x)
 
     def cfun(x):
-      return bc.transform.switch(x, branches, x)
+      return bst.transform.switch(x, branches, x)
 
     self.assertEqual(fun(-1), cfun(-1))
     self.assertEqual(fun(0), cfun(0))
@@ -89,7 +89,7 @@ class TestSwitch(unittest.TestCase):
       return branches[i](x, x)
 
     def cfun(x):
-      return bc.transform.switch(x, branches, x, x)
+      return bst.transform.switch(x, branches, x, x)
 
     self.assertEqual(fun(-1), cfun(-1))
     self.assertEqual(fun(0), cfun(0))
@@ -122,13 +122,13 @@ class TestSwitch(unittest.TestCase):
     branches3 = branches2 + [lambda x: jnp.sin(x) + jnp.cos(x)]  # requires one more residual slot
 
     def fun1(x, i):
-      return bc.transform.switch(i + 1, branches1, x)
+      return bst.transform.switch(i + 1, branches1, x)
 
     def fun2(x, i):
-      return bc.transform.switch(i + 1, branches2, x)
+      return bst.transform.switch(i + 1, branches2, x)
 
     def fun3(x, i):
-      return bc.transform.switch(i + 1, branches3, x)
+      return bst.transform.switch(i + 1, branches3, x)
 
     fwd1, bwd1 = get_conds(fun1)
     fwd2, bwd2 = get_conds(fun2)
@@ -148,7 +148,7 @@ class TestSwitch(unittest.TestCase):
 
   def testOneBranchSwitch(self):
     branch = lambda x: -x
-    f = lambda i, x: bc.transform.switch(i, [branch], 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))
@@ -166,12 +166,12 @@ class TestSwitch(unittest.TestCase):
 class TestIfElse(unittest.TestCase):
   def test1(self):
     def f(a):
-      return bc.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,
+      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,
@@ -183,38 +183,38 @@ class TestIfElse(unittest.TestCase):
 
   def test_vmap(self):
     def f(operands):
-      f = lambda a: bc.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,
+      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)
+                                         a)
       return jax.vmap(f)(operands)
 
-    r = f(bc.random.randint(-20, 20, 200))
+    r = f(bst.random.randint(-20, 20, 200))
     self.assertTrue(r.size == 200)
 
   def test_grad1(self):
     def F2(x):
-      return bc.transform.ifelse((x >= 10, x < 10),
-                                 [lambda x: x, lambda x: x ** 2, ],
-                                 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 bc.transform.ifelse((x >= 10, jnp.logical_and(x >= 0, x < 10), x < 0),
-                                 [lambda x: 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)
+                                  x)
 
     self.assertTrue(jax.grad(F3)(9.0) == 18.)
     self.assertTrue(jax.grad(F3)(11.0) == 1.)