brainstate.transform.ifelse

Contents

brainstate.transform.ifelse#

brainstate.transform.ifelse(conditions, branches, *operands, check_cond=True)#

Represent multi-way if/elif/else control flow.

Parameters:

  • conditions (Sequence) – Sequence of mutually exclusive boolean predicates. When check_cond is True, exactly one entry must evaluate to True.

  • branches (Sequence[Callable]) – Sequence of branch callables evaluated lazily. Must have the same length as conditions, contain at least two callables, and each branch receives *operands when selected.

  • *operands (Any) – Operands forwarded to the selected branch as positional arguments.

  • check_cond (bool) – Whether to verify that exactly one condition evaluates to True.

Returns:

Value produced by the branch corresponding to the active condition.

Return type:

Any

Notes

Unlike a Python if/elif chain, conditions are evaluated together rather than first-true-wins, so when check_cond is True (the default) they must be mutually exclusive: exactly one condition may evaluate to True, otherwise a runtime error is raised. With check_cond=False the check is skipped and the first True condition wins (falling back to the last branch when none is True).

Examples

>>> import jax.numpy as jnp
>>> import brainstate
>>>
>>> def grade(a):
...     return brainstate.transform.ifelse(
...         conditions=[a > 5, jnp.logical_and(a > 0, a <= 5), a <= 0],
...         branches=[
...             lambda: 2.0,  # greater than five
...             lambda: 1.0,  # positive
...             lambda: 0.0,  # non-positive
...         ],
...     )
>>>
>>> grade(jnp.asarray(7.0))   # 2.0
>>> grade(jnp.asarray(-1.0))  # 0.0
Contents