Skip to content

← JAX Vectorization & Control Flow step 16 of 25

Medium Primitives

Multi-Branch Dispatch via lax.switch

Jump to the task

Choose one operation for an array

lax.switch(index, branches, operand) selects a branch by integer index and passes operand to it. The branch list is fixed when the function is traced, while the index can be a runtime value. This provides more than the two choices available with lax.cond.

Here the three choices are sine, cosine and the natural exponential:

import jax
import jax.numpy as jnp
from jax import lax

@jax.jit
def apply_op(x, index):
    branches = [jnp.sin, jnp.cos, jnp.exp]
    return lax.switch(jnp.int32(index), branches, x)

x = jnp.array([0.0, 1.0])
apply_op(x, 0.0)  # approximately [0.0, 0.841471]
apply_op(x, 1.0)  # approximately [1.0, 0.540302]
apply_op(x, 2.0)  # approximately [1.0, 2.718282]

One index selects an operation for the whole array, not a separate operation for each element. Sine and cosine use radians. The exponential is e raised to each input value, not 10 raised to that value.

Pass functions, not computed arrays

The branch list contains callables such as jnp.sin, not values such as jnp.sin(x). Each branch receives the same operands and must return the same structure, shapes and dtypes. All branches are traced, so every branch must be valid even if a particular call selects only one.

Scalar switch expresses conditional execution, subject to compiler transformations. Batching the index with vmap can turn that decision into selection and evaluate multiple branch computations. Do not use Python side effects to infer which branch ran.

A Python if/elif chain controlled by a traced runtime index instead asks Python for a boolean during tracing and raises TracerBoolConversionError. Static Python choices can still work; switch is useful when the choice needs to remain dynamic.

Know what an out-of-range index means

Switch clamps to the nearest valid branch. With three branches, index -1 selects branch 0, not the last branch as Python list indexing would. Index 3 selects branch 2. This is documented behavior, not an input-validation error. Clipping to [0, 2] again does not reject an invalid input; an application that needs rejection must validate it.

This exercise supplies only indices 0, 1 and 2, possibly represented as floats. Cast with jnp.int32(op_index). That representation is part of the exercise interface, not something JAX requires for indices. Fractional indices are outside the contract; casting would truncate them.

Problem

Implement switch_op(x, op_index) using lax.switch with branches [jnp.sin, jnp.cos, jnp.exp] in that order. The input x is a one-dimensional real floating-point JAX array; op_index is a scalar whose value is 0, 1 or 2.

Return the selected operation's array with the same shape and dtype as x. An empty input returns an empty array. The output tests verify the selected values but cannot enforce the use of switch over another equivalent implementation.