Skip to content

← JAX Profiling & Debugging step 10 of 12

Medium Primitives

jaxpr with Multiple Args

Jump to the task

Why this matters

A jaxpr describes inputs as a list of variables called invars. Each ordinary array argument contributes one input variable, regardless of how many elements that array contains.

Structured arguments are flattened first. A dictionary with two array leaves can contribute two input variables even though the Python function accepts just one dictionary.

Worked mini-example

import jax
import jax.numpy as jnp

x = jnp.ones(3)
y = jnp.ones(3)
trace = jax.make_jaxpr(lambda x, y: x + y)(x, y)
count = len(trace.jaxpr.invars)  # 2

def add_pair(pair):
    return pair["left"] + pair["right"]

tree_trace = jax.make_jaxpr(add_pair)({"left": x, "right": y})
tree_count = len(tree_trace.jaxpr.invars)  # also 2

The first function has two array parameters. The second has one dictionary parameter with two array leaves. Both traces have two input variables.

Common pitfalls

  • Count variables, not array elements. Two arrays of length 100 still contribute two variables in this example.
  • Static arguments selected with make_jaxpr's static_argnums are used during tracing; they do not become ordinary input variables.
  • constvars represents separately captured constants, while invars represents dynamic inputs. Scalar literals can appear directly in equations.
  • JAX's object wrappers have changed across releases. This exercise uses trace.jaxpr.invars; do not depend on an exact wrapper class name or the names of printed variables.

Problem

Implement jaxpr_invar_count(x, y). Both inputs are one-dimensional JAX arrays with the same shape.

Trace exactly lambda x, y: x + y with jax.make_jaxpr, access .jaxpr.invars, and return its length as a jnp.float32 scalar.

The result is 2.0, including when both arrays are empty. Use introspection to practice the API. The numeric tests check the result; they cannot distinguish introspection from a constant answer.