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.