Need an operation that goes beyond the familiar NumPy-style interface? jax.lax is JAX’s lower-level API: jax.numpy expresses its operations through jax.lax primitives, many of which wrap equivalent XLA operations. Prefer jax.numpy when it provides what you need because its interface is more stable. section 3.4
jax.lax groups primitives for several kinds of work:
- Structured control flow: lax.cond, lax.switch, lax.while_loop, lax.fori_loop, lax.scan, lax.associative_scan, and lax.map. section 3.4.1
- General mathematical operations, including conv_general_dilated, which provides a more general convolution primitive than jnp.convolve. section 3.4
- Gradient, parallelism, and linear algebra operators, with further treatment deferred to later use. section 3.4
- Explicit type handling: mixed data types generally require manual promotion. section 3.4.2
Choose an image augmentation with lax.switch
Create a list of augmentation functions, draw an integer index with jax.random, and pass both to lax.switch: augmented_image = lax.switch(augmentation_index, augmentations, image). The selected branch applies to the image. Unlike a large Python loop or branch expansion, these structured primitives are JIT-able and differentiable. section 3.4.1
Promotion must be explicit
jnp.add(42, 42.0) performs automatic promotion, but lax.add(42, 42.0) raises a dtype mismatch error. Convert explicitly, for example with lax.add(jnp.float32(42), 42.0). section 3.4.2
Figure 3.9 Diagram of the JAX API layers
Use the API-layer diagram to see how the higher-level jax.numpy interface is built over lower-level jax.lax primitives. section 3.4
