jax.lax.switch
jax.lax.switch#
- jax.lax.switch(index, branches, *operands, operand=<object object>)[source]#
Apply exactly one of
branches
given byindex
.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)
- Parameters
- Returns
Value (B) of
branch(*operands)
for the branch that was selected based onindex
.