Skip to content

← JAX Conceptual Deep-Dives step 11 of 15

Medium Primitives

Counting Jaxpr Equations

Jump to the task

Count operations in a trace

jax.make_jaxpr(f)(x) describes the JAX operations traced for f with input x's shape and dtype. Its .eqns sequence contains the equations at the top level of that trace. Counting them is a way to inspect its structure, not to measure how expensive the calculation will be.

A single equation can operate on a large array or contain another jaxpr. A small top-level count therefore does not imply little work. It also does not show whether the compiler fused several operations into a device kernel: that happens at a different stage.

Worked example

import jax
import jax.numpy as jnp

x = jnp.ones(3, dtype=jnp.float32)
trace = jax.make_jaxpr(lambda x: jnp.sum(x ** 2 + jnp.exp(x)))(x)
operations = [equation.primitive.name for equation in trace.eqns]
count = len(trace.eqns)
print(operations)
# Commonly: ['integer_pow', 'exp', 'add', 'reduce_sum']

The expression squares each element, exponentiates each element, adds the two arrays, then sums the result. In the shown representation these are four equations, despite fitting in one Python expression. Primitive names and decomposition can change with JAX releases; inspect your own trace.

Problem

Implement count_jaxpr_eqns(x):

  1. Trace exactly lambda x: jnp.sum(x ** 2 + jnp.exp(x)) on x.
  2. Read .eqns on the returned object.
  3. Return jnp.float32(len(trace.eqns)).

x is a one-dimensional float32 JAX array. Return a scalar JAX value with dtype float32. Count only the top-level equations; do not recursively count nested traces, count printed lines, or return the expression's numerical sum.

The checker computes the count using the same JAX installation as your solution, so it does not assume a fixed count for every release. It also checks the return shape and dtype. Derive the count from the trace rather than returning the number shown above.

What the trace tells you

Tracing executes Python with abstract array values to record operations; it is not a timed numerical evaluation of the function. Python side effects can still occur while tracing, so keep the traced function free of them.

For this expression, changing the entries of x does not change its operation sequence. Changing its length changes the type annotations but usually not the equation count. Other functions can branch on shapes or unroll loops based on dimensions, so shape independence is not a general property of jaxprs.

Use the equations to ask what was traced. Use lowered or compiled representations and profiling to investigate optimization or runtime cost.