tf.keras.ops.switch

Apply exactly one of the branches given by index.

tf.keras.ops.switch(
    index, branches, *operands
)

If index is out of bounds, it is clamped to within bounds.

The semantics of switch are given roughly by this Python implementation:

def switch(index, branches, *operands):
    index = clamp(0, index, len(branches) - 1)
    return branches[index](*operands)

index An integer scalar indicating which branch function to apply.
branches A sequence of functions to be applied based on index.
operands Inputs to whichever branch is applied.

The outputs of branch(*operands) for the branch that was selected based on index.

Examples:

add_fn = lambda x, y: x + y
subtract_fn = lambda x, y: x - y
x = keras.ops.array(2.0)
y = keras.ops.array(0.5)
branches = [add_fn, subtract_fn]
keras.ops.switch(0, branches, x, y)
2.5
 
keras.ops.switch(1, branches, x, y)
1.5