Shape, dtype, and no numbers

Put a print(type(x)) at the top of a jitted function and the class that comes back is DynamicJaxprTracer. The chapter above says JAX calls your function with stand-ins; this is the stand-in, named, and the print fires at trace time for the reason that chapter already gave.

Ask the tracer for its aval and it answers ShapedArray(float32[3]). That object is an abstract value: a shape, a dtype, and a couple of flags, with no storage behind it. Two arrays holding completely different numbers at the same shape and dtype have the same aval, so one recording serves both.

Plenty of what you would ask a real array still works, and it is worth being exact about why. x.shape gives a tuple of ordinary Python ints, len(x) gives an int, x.dtype gives a dtype. Each of those answers was already sitting in the aval, so nothing had to be computed and nothing had to exist.

The last line of the run is the one to carry into the next lesson. x.sum() > 0 inside the trace has the aval ShapedArray(bool[]), and bool[] is exactly the phrase the branch error prints back at you. That message is not describing your data. It is printing the aval.

import jax
import jax.numpy as jnp

@jax.jit
def probe(x):
    print(type(x).__name__)
    print(x.aval, x.aval == jax.core.ShapedArray((3,), jnp.float32))
    print(x.shape, x.dtype, x.ndim, len(x))
    print((x.sum() > 0).aval)
    print(isinstance(x, jax.Array))
    return x

probe(jnp.ones(3))

# DynamicJaxprTracer
# ShapedArray(float32[3]) True
# (3,) float32 1 3
# ShapedArray(bool[])
# True

Everything the trace can reason about a value is in its aval, and the aval is small. A shape as a tuple of ints, a dtype, and a weak_type flag that records whether the value came from a Python scalar rather than an array. Nothing else about the value survives into the recording.

That weak_type flag is where the promotion rule from chapter 1 is stored. Pass the Python float 2.0 into a jitted function and its aval prints as ShapedArray(float32[], weak_type=True); pass jnp.float32(2.0) and the flag is gone. Chapter 1 owns what the flag does to a promotion, and chapter 3 owns what it does to a cache key. Here it is one more field of an abstract value, and it is why two calls that look identical in Python can present different avals.

Comparing avals is the cheapest test of the mental model there is. x.aval == jax.core.ShapedArray((3,), jnp.float32) answers True in the run above, which means you can write down what the trace sees before you run anything, then check yourself in one line.

If the trace only ever sees shapes and dtypes, you can run one without owning the arrays. jax.eval_shape(f, *args) does exactly that: it traces f, computes the output avals, and allocates nothing, so the arguments can be ShapeDtypeStruct descriptions instead of real data.

The use for that is practical rather than decorative. A model whose parameters would not fit on the machine in front of you still answers shape questions in milliseconds, and a shape bug that would have surfaced after a long allocation surfaces before it. No compile happens either, which chapter 3 measures the cost of.

One thing does carry over, and it is the subject of the next lesson. eval_shape runs your Python once, exactly as jit does, so a function that refuses to trace refuses here too. Cheap tracing is still tracing.

import jax
import jax.numpy as jnp
from jax import ShapeDtypeStruct

def block(x, w):
    return jnp.tanh(x @ w).sum(axis=1)

print(jax.eval_shape(block, ShapeDtypeStruct((8, 16), jnp.float32),
                            ShapeDtypeStruct((16, 4), jnp.float32)))
print(jax.eval_shape(lambda x: x, 2.0))
print(jax.eval_shape(lambda x: x, jnp.float32(2.0)))

# ShapeDtypeStruct(shape=(8,), dtype=float32)
# ShapeDtypeStruct(shape=(), dtype=float32, weak_type=True)
# ShapeDtypeStruct(shape=(), dtype=float32)

jit is not the only thing that traces, and the three transforms do not send the same object through your function. Under jit you get a DynamicJaxprTracer, under vmap a BatchTracer, under grad a JVPTracer. Same interface, three different jobs.

Look at what vmap reports and the design shows through. The input was (2, 3) and the aval inside the function is float32[3], one row, because vmap hides the mapped axis from the code it wraps. Your function is written for one example and the batch axis is bookkeeping that happens around it.

Stack all three and the stand-ins nest in the order the transforms were applied. jit(vmap(grad(f))) hands your function a JVPTracer whose primal is a BatchTracer whose value is a DynamicJaxprTracer, and only that innermost one still carries the full float32[2,3]. Unwrapping the stack by hand is a two-line trick and the fastest way to see which transform is currently in charge.

import jax
import jax.numpy as jnp

def who(x):
    print(type(x).__name__, x.aval)
    print(type(x.primal).__name__, x.primal.aval)
    inner = x.primal.val
    print(type(inner).__name__, inner.aval)
    return jnp.sum(x)

jax.jit(jax.vmap(jax.grad(who)))(jnp.ones((2, 3)))

# JVPTracer ShapedArray(float32[3])
# BatchTracer ShapedArray(float32[3])
# DynamicJaxprTracer ShapedArray(float32[2,3])