brainstate 0.0.2.post20240825__py2.py3-none-any.whl → 0.0.2.post20240910__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
@@ -1597,6 +1597,6 @@ def _get_delay(delay_time, delay_step):
       delay_time = delay_step * environ.get_dt()
   else:
     assert delay_step is None, '"delay_step" should be None if "delay_time" is given.'
-    assert isinstance(delay_time, (int, float))
+    # assert isinstance(delay_time, (int, float))
     delay_step = math.ceil(delay_time / environ.get_dt())
   return delay_time, delay_step

brainstate/_state.py

@@ -234,6 +234,8 @@ class ShapeDtype:
   def __init__(self, shape, dtype):
     self.shape = shape
     self.dtype = dtype
+    self.ndim = len(shape)
+    self.size = np.prod(shape)
 
   def __repr__(self):
     return f'{self.dtype}{list(self.shape)}'

brainstate/nn/_projection/_align_post.py

@@ -15,6 +15,8 @@
 
 from typing import Optional, Union
 
+
+import brainunit as u
 from brainstate._module import (register_delay_of_target,
                                 Projection,
                                 Module,
@@ -278,11 +280,19 @@ class FullProjAlignPostMg(Projection):
     self.comm = comm
 
     # delay initialization
-    if delay is not None and delay > 0.:
-      delay_cls = register_delay_of_target(pre)
-      delay_cls.register_entry(self.name, delay)
-      self.delay = delay_cls
-      self.has_delay = True
+    if delay is not None:
+      if isinstance(delay, u.Quantity):
+        has_delay = delay.mantissa > 0.
+      else:
+        has_delay = delay > 0.
+      if has_delay:
+        delay_cls = register_delay_of_target(pre)
+        delay_cls.register_entry(self.name, delay)
+        self.delay = delay_cls
+        self.has_delay = True
+      else:
+        self.delay = None
+        self.has_delay = False
     else:
       self.delay = None
       self.has_delay = False
@@ -502,11 +512,19 @@ class FullProjAlignPost(Projection):
     self.out = out
 
     # delay initialization
-    if delay is not None and delay > 0.:
-      delay_cls = register_delay_of_target(pre)
-      delay_cls.register_entry(self.name, delay)
-      self.delay = delay_cls
-      self.has_delay = True
+    if delay is not None:
+      if isinstance(delay, u.Quantity):
+        has_delay = delay.mantissa > 0.
+      else:
+        has_delay = delay > 0.
+      if has_delay:
+        delay_cls = register_delay_of_target(pre)
+        delay_cls.register_entry(self.name, delay)
+        self.delay = delay_cls
+        self.has_delay = True
+      else:
+        self.delay = None
+        self.has_delay = False
     else:
       self.delay = None
       self.has_delay = False

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/surrogate.py

@@ -158,6 +158,9 @@ class Sigmoid(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def sigmoid(
     x: jax.Array,
@@ -243,6 +246,9 @@ class PiecewiseQuadratic(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def piecewise_quadratic(
     x: jax.Array,
@@ -339,6 +345,9 @@ class PiecewiseExp(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def piecewise_exp(
     x: jax.Array,
@@ -426,6 +435,9 @@ class SoftSign(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def soft_sign(
     x: jax.Array,
@@ -508,6 +520,9 @@ class Arctan(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def arctan(
     x: jax.Array,
@@ -589,6 +604,9 @@ class NonzeroSignLog(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def nonzero_sign_log(
     x: jax.Array,
@@ -683,6 +701,9 @@ class ERF(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def erf(
     x: jax.Array,
@@ -780,6 +801,9 @@ class PiecewiseLeakyRelu(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(c={self.c}, w={self.w})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.c, self.w))
+
 
 def piecewise_leaky_relu(
     x: jax.Array,
@@ -898,6 +922,9 @@ class SquarewaveFourierSeries(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(n={self.n}, t_period={self.t_period})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.n, self.t_period))
+
 
 def squarewave_fourier_series(
     x: jax.Array,
@@ -988,6 +1015,9 @@ class S2NN(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta}, epsilon={self.epsilon})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.beta, self.epsilon))
+
 
 def s2nn(
     x: jax.Array,
@@ -1089,6 +1119,9 @@ class QPseudoSpike(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def q_pseudo_spike(
     x: jax.Array,
@@ -1178,6 +1211,9 @@ class LeakyRelu(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.beta))
+
 
 def leaky_relu(
     x: jax.Array,
@@ -1277,6 +1313,9 @@ class LogTailedRelu(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def log_tailed_relu(
     x: jax.Array,
@@ -1368,6 +1407,9 @@ class ReluGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, width={self.width})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.width))
+
 
 def relu_grad(
     x: jax.Array,
@@ -1446,6 +1488,9 @@ class GaussianGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, sigma={self.sigma})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.sigma))
+
 
 def gaussian_grad(
     x: jax.Array,
@@ -1530,6 +1575,9 @@ class MultiGaussianGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(h={self.h}, s={self.s}, sigma={self.sigma}, scale={self.scale})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.h, self.s, self.sigma, self.scale))
+
 
 def multi_gaussian_grad(
     x: jax.Array,
@@ -1615,6 +1663,9 @@ class InvSquareGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def inv_square_grad(
     x: jax.Array,
@@ -1685,6 +1736,9 @@ class SlayerGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def slayer_grad(
     x: jax.Array,

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)