Skip to content

← JAX Foundations step 2 of 25

Medium End-to-End

Tracing: Shape vs Value

Jump to the task

Why this matters

In ordinary JIT tracing, JAX knows an input array's shape and dtype, but not its values. That distinction determines which conditions Python can evaluate while the function is being traced.

For a one-dimensional input, if x.shape[0] >= 3: checks a known length. Python chooses that branch during tracing. A different input length can produce another compiled version with a different branch.

In contrast, if jnp.sum(x) > threshold: needs array values. Python cannot turn that traced comparison into a boolean, so it raises TracerBoolConversionError. Use jax.lax.cond for this choice: its scalar predicate can depend on the input values.

Worked mini-example

import jax
import jax.numpy as jnp

def choose_scale(x):
    return jax.lax.cond(
        jnp.sum(x) > 0,
        lambda: x * 2,
        lambda: -x,
    )

choose_jit = jax.jit(choose_scale)
a = choose_jit(jnp.array([1.0, 2.0]))    # [2, 4]
b = choose_jit(jnp.array([-1.0, -2.0]))  # [1, 2]

These inputs have the same shape and dtype. The compiled function can choose different branches from their values without retracing.

Common pitfalls

  • Both lax.cond branches are traced. Keep them pure: do not use Python side effects to detect which branch executes.
  • Corresponding branch outputs must have the same shape and dtype, and the same container structure. Broadcasting compatibility is not enough; outputs of shape (1,) and (3,) do not match.
  • For a scalar predicate, cond represents a choice of one branch, subject to compiler optimizations. Under vmap over predicates, it becomes a selection and may evaluate both branch computations.
  • The shape check here assumes ordinary JIT with concrete input dimensions. Exporting with symbolic dimensions has additional rules; a symbolic comparison may not be decidable.

Problem

Implement shape_or_value_branch(x, threshold). The input x is a one-dimensional floating-point JAX array, and threshold is a float scalar. Always return an array of shape (3,).

If x has at least three elements, use lax.cond to choose:

  • When jnp.sum(x) > threshold, return x[:3] * 2.0.
  • Otherwise, including equality, return x[:3] + 1.0.

If x has fewer than three elements, use the outer Python shape check to return jnp.zeros(3) + threshold. This also covers empty inputs.

The predicate uses the sum of all elements, even though the output uses only the first three. For x = [1, 2, 3, 4, 5] and threshold = 10, the sum is 15, so return [2, 4, 6]. For x = [0.5, 0.5] and threshold = 1, return [1, 1, 1].