Skip to content

← JAX Performance & Distributed step 19 of 25

Medium Primitives

Elementwise Chains and Fusion

Jump to the task

Why this matters

A compiled function gives JAX's compiler a view of several operations together. It can sometimes combine them so that intermediate results do not need separate memory reads and writes. This optimization is called fusion.

Python variable names do not force intermediate arrays to be materialized. These two functions express the same calculation and give the compiler the same opportunity to optimize when each whole function is compiled:

import jax
import jax.numpy as jnp

@jax.jit
def compact(x):
    return ((x * 2.0) + 1.0) ** 2

@jax.jit
def named_steps(x):
    doubled = x * 2.0
    shifted = doubled + 1.0
    return shifted ** 2

Choose the form that makes the calculation easier to read. Neither spelling guarantees one kernel or identical timing on every backend. Fusion decisions depend on the compiler, operations, shapes and hardware.

Worked mini-example

x = jnp.array([-0.5, 0.0, 1.0], dtype=jnp.float32)
result = compact(x)  # [0, 1, 9]
# Double: [-1, 0, 2]; add one: [0, 1, 3]; square: [0, 1, 9].

Problem

Implement fused_chain(x) to return ((x * 2.0) + 1.0) ** 2. Use the direct expression to practice expressing an elementwise chain.

x is a one-dimensional finite float32 JAX array. Assume intermediate values and the result stay finite. Return a float32 array with the same shape, including for empty input. Apply the formula independently to each element; do not reduce or reorder the values.

The fixtures check numerical output. They do not prove that your function was compiled, fused into one kernel, or faster than another implementation. The function name is retained for this exercise; it is not a performance guarantee.

Common pitfalls

  • Parentheses matter: square the shifted value, not just x or the constant.
  • Keep the whole chain inside the compiled function when exploring fusion. Separate eager dispatches do not give the compiler the same view across the chain. JIT is one compilation entry point, not the only setting in which compiler optimization can occur.
  • For the float32 inputs here, Python literals 2 and 2.0 both work. Literal choice is not a general way to control output dtype; specify and check dtypes when translating to other input types.
  • Benchmark after a warm-up call and wait for output readiness with block_until_ready(). Otherwise timing can include compilation or measure dispatch without waiting for device work. Inspect compiled code or a profiler to investigate fusion; correct values alone cannot show it.