TopolearnXLA can significantly speed up compute-heavy models and is required for TPUs.
XLA (Accelerated Linear Algebra) is a compiler that takes a whole TensorFlow graph and produces optimised machine code for CPU, GPU or TPU. Its main win is fusion: chains of element-wise ops that TF would launch as separate kernels, each writing to memory, become one kernel.
Enable it per function:
@tf.function(jit_compile=True)
def train_step(x, y): ...or for Keras with model.compile(..., jit_compile=True). In Keras 3 the default jit_compile="auto" already enables XLA where the model supports it.
Constraints and costs:
- XLA compiles for static shapes. Each new input shape triggers a new compilation, so pad or bucket variable-length inputs.
- Some ops have no XLA kernel; compilation then fails with an error naming the op. Ops with data-dependent output shapes (such as
tf.unique) are a common cause. - The first call pays compilation time.
Inspect what XLA produces with f.experimental_get_compiler_ir(*args)(stage="hlo"). On TPUs everything is compiled through XLA, so the same constraints always apply there. XLA is also the compiler behind JAX, so the skills transfer.
Example
import tensorflow as tf
@tf.function(jit_compile=True)
def dense_gelu(x, w, b):
y = tf.matmul(x, w) + b
# tanh-approximate GELU: several element-wise ops XLA fuses into one kernel
return 0.5 * y * (1.0 + tf.tanh(0.7978845608 * (y + 0.044715 * y ** 3)))
x = tf.random.normal([256, 512])
w = tf.random.normal([512, 512])
b = tf.zeros([512])
print(dense_gelu(x, w, b).shape) # (256, 512)
hlo = dense_gelu.experimental_get_compiler_ir(x, w, b)(stage="hlo")
print(hlo.splitlines()[0]) # HloModule ... (the compiled program)
dense_gelu(tf.random.normal([128, 512]), w, b) # new shape -> recompilesRemember
jit_compile=Truecompiles a tf.function with XLA- Fusion of element-wise ops is the main speedup
- Static shapes: new shapes recompile; some ops are unsupported
Watch out for
- Feeding many different sequence lengths to an XLA function, recompiling constantly.