How can JAX compile a function for many inputs without knowing their runtime values? During tracing, the default ShapedArray tracer records each input's shape and data type, but not its concrete value. JIT compilation can therefore reuse one compiled function for different values with the same shape and type. section 5.2.1
Tracing with ShapedArray has practical consequences:
- Shape-dependent loops such as range(x.shape[0]) can be traced, while loops whose bounds use x's value raise TracerIntegerConversionError. section 5.2.1
- Value-dependent Python conditionals such as if x > 0 raise TracerBoolConversionError; use static arguments or jax.lax.cond instead. section 5.2.1
- Python side effects occur during tracing but do not appear in the resulting Jaxpr, while global values may be captured as constants. section 5.2.1
Listing 5.10 Using jax.make_jaxpr()
def f1(x, y, z): #1
return jnp.sum(x + y * z)
x = jnp.array([1.0, 1.0, 1.0]) #2
y = jnp.ones((3,3))*2.0
z = jnp.array([2.0, 1.0, 0.0]).T
jax.make_jaxpr(f1)(x,y,z) #3
>>> { lambda ; a:f32[3] b:f32[3,3] c:f32[3]. let
>>> d:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] c
>>> e:f32[3,3] = mul b d
>>> f:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] a
>>> g:f32[3,3] = add f e
>>> h:f32[] = reduce_sum[axes=(0, 1)] g
>>> in (h,) }
Inspecting make_jaxpr output makes abstract inputs visible through entries such as f32[3] and f32[3,3], alongside the operations retained for compilation. section 5.2.1
