Jax

jax Best Practices

JAX is the backbone of our AI/ML and numerical computing projects. Adhere to these principles for high-performance, reproducible, and maintainable JAX code.

1. Functional Purity: The Absolute Core

JAX transformations (jit, grad, vmap, pmap) operate exclusively on functionally pure code. This means functions must be free of side-effects: all inputs explicit, all results returned.

• Avoid mutable global state: JAX captures global values at first jit compilation, leading to stale values. ❌ BAD:

  g = 0
  def impure_uses_globals(x):
      return x + g # `g` is captured at first jit
  # ... later g = 10, but jit(impure_uses_globals) still uses g=0
  

  ✅ GOOD: Pass all state explicitly.

  def pure_uses_globals(x, g_val):
      return x + g_val
  # ... pass g_val=0, then g_val=10
  

• No in-place array mutation: JAX arrays are immutable. Use the .at[] syntax for functional updates. ❌ BAD:

  import jax.numpy as jnp
  arr = jnp.zeros((3,3))
  arr[1, :] = 1.0 # TypeError!
  

  ✅ GOOD:

  import jax.numpy as jnp
  arr = jnp.zeros((3,3))
  updated_arr = arr.at[1, :].set(1.0) # Returns a new array
  

• Avoid Python iterators in jitted code: Iterators introduce state. ❌ BAD:

  from jax import jit
  def sum_iterator(it):
      total = 0
      for x in it: # Python loop with iterator
          total += x
      return total
  # jit(sum_iterator)(iter(range(10))) # Will fail or give unexpected results
  

  ✅ GOOD: Use JAX control flow primitives.

  from jax import lax
  import jax.numpy as jnp
  def sum_array(arr):
      # Use lax.scan or lax.fori_loop for JAX-compatible loops
      return lax.fori_loop(0, arr.shape[0], lambda i, x: x + arr[i], 0)
  # jit(sum_array)(jnp.arange(10))
  

2. Numerical Type Discipline

Prioritize float32 for performance on accelerators. Avoid implicit float64 promotion.

• Explicit dtype for constants: ❌ BAD: Implicitly typed Python floats can lead to float64 promotion.

  import jax.numpy as jnp
  x = jnp.ones(5, dtype=jnp.float32)
  y = x * 2.0 # 2.0 is a Python float, can cause unwanted promotion
  

  ✅ GOOD: Use 0-D jnp.array with explicit dtype or jnp.float32().

  import jax.numpy as jnp
  x = jnp.ones(5, dtype=jnp.float32)
  y = x * jnp.array(2.0, dtype=jnp.float32)
  # Or more concisely:
  z = x * jnp.float32(2.0)
  

3. Performance Considerations: Stable Compilation

Prevent costly recompilations and leverage XLA effectively.

• Static shapes: Keep input shapes static or pass them via static_argnums to jit. Dynamic shapes force recompilation. ❌ BAD:

  from jax import jit
  def dynamic_shape_func(x):
      return x.sum()
  # jit(dynamic_shape_func)(jnp.ones(5))
  # jit(dynamic_shape_func)(jnp.ones(10)) # Recompiles!
  

  ✅ GOOD:

  from jax import jit
  @jit
  def static_shape_func(x):
      return x.sum()
  # Call with consistent shapes
  # static_shape_func(jnp.ones(5))
  # static_shape_func(jnp.ones(5)) # Uses cached compilation
  

  If shapes must vary, consider static_argnums for non-array arguments that determine shape.

• JAX control flow primitives: Always use lax.scan, lax.while_loop, lax.cond inside jitted functions. Python control flow breaks XLA compilation. ❌ BAD:

  from jax import jit
  @jit
  def python_loop(x, n):
      for _ in range(n): # Python loop inside jit
          x = x * 2
      return x
  

  ✅ GOOD:

  from jax import jit, lax
  @jit
  def jax_loop(x, n):
      # lax.fori_loop(start, stop, body_fn, init_val)
      return lax.fori_loop(0, n, lambda i, val: val * 2, x)
  

• Vectorization with vmap: For data-parallel batching, vmap is your tool.

  from jax import vmap
  import jax.numpy as jnp
  def elementwise_op(x, y):
      return x + y # Operates on scalars or single elements
  batched_op = vmap(elementwise_op)
  # batched_op(jnp.array([1,2,3]), jnp.array([4,5,6])) # Applies elementwise
  

4. Code Organization and Randomness

• Standard Imports:

  import jax
  import jax.numpy as jnp
  import jax.random as jr
  import jax.lax as lax
  

• Randomness Management: Use jax.random and explicitly split PRNGKeys. Never reuse a key. ❌ BAD:

  import jax.random as jr
  key = jr.PRNGKey(0)
  val1 = jr.normal(key, (5,)) # Key used
  val2 = jr.normal(key, (5,)) # Key reused, not independent!
  

  ✅ GOOD:

  import jax.random as jr
  key = jr.PRNGKey(0)
  key, subkey1 = jr.split(key)
  val1 = jr.normal(subkey1, (5,))
  key, subkey2 = jr.split(key)
  val2 = jr.normal(subkey2, (5,)) # Independent random values
  

5. Type Hints

Use standard Python type hints, especially jax.Array for JAX arrays. This improves readability and enables static analysis.

import jax.numpy as jnp
from jax import Array, jit

@jit
def add_arrays(a: Array, b: Array) -> Array:
    """Adds two JAX arrays."""
    return a + b

# Example usage:
# result = add_arrays(jnp.ones(5), jnp.ones(5))

6. Testing Approaches

• pytest is standard: Use pytest for unit and integration tests.
• Gradient checking: For custom operations or complex functions, use jax.test_util.check_grads to verify gradients.

import jax
import jax.numpy as jnp
from jax.test_util import check_grads

def my_complex_func(x):
    return jnp.sin(x) * jnp.exp(x)

# Test gradients numerically vs. analytically
# check_grads(my_complex_func, (jnp.array(1.0),), order=1)