Skip to content

← JAX Performance & Distributed step 24 of 25

Medium Primitives

Shape Polymorphism (Concept)

Jump to the task

Why this matters

A function can work for several input shapes without using symbolic dimensions. For example, jnp.sum(x) handles vectors of different lengths because the source code does not assume a particular length.

Ordinary jax.jit usually specializes compiled code to concrete shapes and dtypes. Calling the same function with length three and then length five can create separate specializations. Returning to length three can reuse its cached specialization; changing the values alone normally does not require a new one.

Symbolic shape export addresses a different stage. It lets JAX trace and lower a function once using dimension expressions, then reuse that exported representation for compatible concrete shapes. It does not promise one machine-code binary that runs every shape without further compilation.

Worked mini-example

import jax
import jax.numpy as jnp
from jax import export

def sum_vector(x):
    return jnp.sum(x)

b, = export.symbolic_shape("b")
specification = jax.ShapeDtypeStruct((b,), jnp.float32)
exported = export.export(jax.jit(sum_vector))(specification)

first = exported.call(jnp.array([1.0, 2.0, 3.0]))  # 6
second = exported.call(jnp.array([10.0, 20.0, 30.0, 40.0, 50.0]))  # 150

symbolic_shape returns a tuple of dimension expressions, so the comma in b, = ... extracts its one dimension. Here b represents a positive integer. Calls must satisfy the exported rank, dtype and shape constraints. The export is reused, while execution may compile for new concrete shapes.

Problem

Implement polymorphic_sum(x) to return the sum of a one-dimensional float32 JAX array as a float32 scalar. Accept any length, including zero; an empty vector sums to zero. Assume inputs and the sum are finite.

You only need to implement the sum, not the export machinery above. The fixtures exercise different vector lengths. They do not verify symbolic export, cache reuse or the number of compilations.

Common pitfalls

  • Shape-independent source code is not automatically a polymorphic executable. Distinguish Python tracing, lowering and machine-code compilation when discussing what work is reused.
  • Shape and dtype are not the entire JIT cache key. Static arguments, function identity and other compilation settings can also matter.
  • A symbolic dimension is not an arbitrary data-dependent size. Exported programs still have shape constraints, and not every operation or shape relationship can be expressed or proved during symbolic tracing.
  • The bare symbolic dimension b above is positive, so that export does not accept length zero even though the ordinary sum function does.
  • Keep the jitted function object when testing cache reuse. Creating a new function repeatedly can defeat the reuse you intended to measure.