brainpy.math.controls.cond#

brainpy.math.controls.cond(pred, true_fun, false_fun, operands, dyn_vars=None)[source]#

Simple conditional statement (if-else) with instance of Variable.

>>> import brainpy.math as bm
>>> a = bm.Variable(bm.zeros(2))
>>> b = bm.Variable(bm.ones(2))
>>> def true_f(_):  a.value += 1
>>> def false_f(_): b.value -= 1
>>>
>>> bm.cond(True, true_f, false_f, dyn_vars=[a, b])
>>> a, b
Variable([1., 1.], dtype=float32), Variable([1., 1.], dtype=float32)
>>>
>>> bm.cond(False, true_f, false_f, dyn_vars=[a, b])
>>> a, b
Variable([1., 1.], dtype=float32), Variable([0., 0.], dtype=float32)
Parameters

  • pred (bool) – Boolean scalar type, indicating which branch function to apply.

  • true_fun (callable, jnp.ndarray, JaxArray, float, int, bool) – Function to be applied if pred is True. This function must receive one arguement for operands.

  • false_fun (callable, jnp.ndarray, JaxArray, float, int, bool) – Function to be applied if pred is False. This function must receive one arguement for operands.

  • operands (Any) – Operands (A) input to branching function depending on pred. The type can be a scalar, array, or any pytree (nested Python tuple/list/dict) thereof.

  • dyn_vars (optional, Variable, sequence of Variable, dict) – The dynamically changed variables.

Returns

res – The conditional results.

Return type

Any