Switch

Functional branch selection.

Abstract Signature:

Switch(index: int, branches: List[Callable])

JAX (Core)

API: jax.lax.switch
Strategy: Direct Mapping
Official Docs ↗

Keras

API: keras.ops.switch
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: keras.ops.switch
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: flax.nnx.switch
Strategy: Direct Mapping
Official Docs ↗