XLA, or Accelerated Linear Algebra, is a domain-specific compiler for linear algebra and array programs. In JAX, jit() sends numerical functions through XLA to produce efficient code for CPUs, GPUs, and TPUs. section 1.2.2 section 5.2.2
Its main advantage is seeing a composed computation rather than optimizing each operation in isolation. XLA can fuse multiplication, addition, and summation into one GPU kernel, avoid intermediate arrays, and reduce memory transfers. section 5.2.2
Compiling a small JAX computation
Define f(x, y, z) = sum(x + y * z), wrap it with jax.jit(f), and call lower(x, y, z) to inspect the generated StableHLO or HLO before compiling. This exposes the path from a JAX function to device-specific code. section 5.2.2
XLA fits into JAX's transformation model:
- jit() compiles a function through XLA for efficient execution. section 1.2.2
- pmap() compiles and replicates a function across XLA devices for parallel execution. section 1.2.2
- JAX is more than a NumPy interface because compilation joins automatic differentiation, vectorization, and parallelization. section 1.2.2 section 1.2
Listing 5.18 Compiling Python code to StableHLO and HLO
def f(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 #2
z = jnp.array([2.0, 1.0, 0.0]).T #2
f_jitted = jax.jit(f) #3
f_lowered = f_jitted.lower(x,y,z) #4
print(f_lowered.as_text())
>>> ... #5
>>> %5 = stablehlo.add %4, %2 : tensor<3x3xf32> #5
>>> %6 = stablehlo.constant dense<0.000000e+00> : tensor<f32> #5
>>> ... #5
f_compiled = f_jitted.lower(x,y,z).compile() #6
print(f_compiled.as_text())
>>> ... #7
>>> %multiply.1 = f32[3,3]{1,0} multiply(f32[3,3]{1,0}
%param_0.4, f32[3,3]{1,0} %broadcast.4),metadata={op_name="jit(f)/jit(main)/mul"
source_file="<ipython-input-99-95d48614b981>"
source_line=2} #7
>>> %add.2 = f32[3,3]{1,0} add(f32[3,3]{1,0}
%broadcast.5, f32[3,3]{1,0} %multiply.1), metadata={op_name="jit(f)/jit(main)/add"
source_file="<ipython-input-99-95d48614b981>" source_line=2} #7
>>> ... #7
#1 A simple function with a few operations
#2 Some test data
#3 JITs the function
#4 Lowers the function (generating StableHLO code)
#5 StableHLO representation
#6 Compiles the lowered function for the specific backend (generating HLO code)
#7 HLO representation
Inspect a complete example that lowers a JAX function to StableHLO and HLO, then compiles it. section 5.2.2
