Tracing Deep Dive: Abstract Values and ShapedArray

During tracing, JAX replaces your real data with abstract values. The most common abstract type is ShapedArray, which knows the array's shape and dtype but not its actual values.

import jax
import jax.numpy as jnp

@jax.jit
def show_tracer(x):
    # During tracing, x is a Tracer wrapping a ShapedArray
    print(f"Type during trace: {type(x)}")
    # Shape is known:
    print(f"Shape: {x.shape}")
    # But values are NOT known:
    # if x[0] > 0: ...  # This would fail!
    return x + 1

show_tracer(jnp.array([1.0, 2.0, 3.0]))
# Prints something like:
# Type during trace: <class 'jax...JaxprTracer'>
# Shape: (3,)

You can inspect what JIT will trace using jax.make_jaxpr, which shows the computation graph (called a jaxpr) without actually compiling it:

import jax
import jax.numpy as jnp

def my_fn(x, y):
    z = jnp.sin(x) + jnp.cos(y)
    return jnp.sum(z)

# See the traced computation
jaxpr = jax.make_jaxpr(my_fn)(jnp.ones(3), jnp.ones(3))
print(jaxpr)
# { lambda ; a:f32[3] b:f32[3]. let
#     c:f32[3] = sin a
#     d:f32[3] = cos b
#     e:f32[3] = add c d
#     f:f32[] = reduce_sum[axes=(0,)] e
#   in (f,) }

This jaxpr is the intermediate representation that gets sent to XLA. Each line is a primitive operation. Notice that Python-level constructs (variable names, loops, function calls) have been traced away — only the math remains.

What Tracing Preserves and What It Loses