Migrating NumPy to JAX: Avoiding Common Pitfalls

1. The Cardinal Sin: Mutating State

The most fundamental difference between NumPy and JAX is immutability. NumPy arrays are mutable containers; you can change their contents at will.

import numpy as np
a_np = np.arange(5.)
a_np[0] = 100. # Perfectly fine
print(a_np) # [100., 1., 2., 3., 4.]

This is so natural in Python that we barely think about it. However, in JAX, this is strictly forbidden. JAX arrays are immutable.

This isn't an arbitrary restriction. Immutability is the bedrock that enables JAX's function transformations. If a function could secretly modify one of its input arrays, it would no longer be pure. The compiler couldn't safely reorder operations, parallelize computations, or reason about derivatives, because the array's value could change unexpectedly.

So, how do we modify arrays? We don't. We create new arrays that incorporate our desired changes. JAX provides a clean, functional syntax for this using the .at property.

import jax.numpy as jnp
a_jax = jnp.arange(5.)




# This would fail: a_jax[0] = 100.




# The functional way:
b_jax = a_jax.at[0].set(100.)

print(f"Original array (unchanged): {a_jax}")
print(f"New array: {b_jax}")

This pattern—new_array = old_array.at[indices].set(values)—is the JAX-idiomatic replacement for in-place assignment. Let's explore this in more depth.

The .at property is versatile, supporting not just .set() but also operations like .add(), .mul(), and .min(), which correspond to NumPy's np.add.at, etc. We will dedicate the next lesson to mastering this API.

2. Tracing vs. Python Control Flow: A Clash of Worlds

The next major pitfall stems from a misunderstanding of how jax.jit and other transformations work. When you apply jit to a function, JAX doesn't just run your Python code faster. It first performs a process called tracing. JAX runs your function once with abstract "tracer" objects instead of concrete numbers. It records every operation performed on these tracers to build a static computation graph (called a jaxpr). This graph is then compiled by XLA into highly optimized machine code.

The critical word here is static. The structure of the computation graph cannot depend on the values of the inputs, only their shapes and dtypes. This is where standard Python control flow often breaks.

Consider this function:

def python_if(x):
  if x > 0: # This condition depends on the VALUE of x
    return x
  else:
    return -x

When Python executes the if statement, it needs a concrete boolean result for x > 0. But during JAX tracing, x is a tracer, not a concrete value. Trying to ask "is this abstract tracer greater than zero?" makes no sense, leading to a ConcretizationTypeError. The same logic applies to Python loops where the loop bounds depend on data, like for i in range(n): where n is a traced array.

The following document, "Crimes Against jax.jit," is a wonderfully concise and practical guide to these issues.

The key insight is that you must use JAX's own control flow primitives, which are designed to work with tracers.

• jnp.where(condition, x, y) is the equivalent of a vectorized if/else. It always evaluates both x and y and selects from them based on the condition, ensuring the computation graph's structure remains static.
• For more complex, non-vectorizable branching, jax.lax.cond is the tool.
• For loops, you must use jax.lax.fori_loop, jax.lax.while_loop, or the powerful jax.lax.scan. We will cover these lax primitives in detail in the next module.

It is crucial to remember that control flow that depends on static information, such as an array's shape (if x.ndim > 1: ...), is perfectly fine. This information is known at trace time and doesn't introduce data dependency into the graph's structure.

3. The General Principle: Side Effects and Purity

Both immutability and control flow issues are specific instances of a more general rule: JAX transformations require pure functions. A pure function's output depends only on its explicit inputs, and it has no observable side effects (like modifying a global variable, printing to the console, or writing to a file).

Let's see what happens when we violate this.

results = []

@jax.jit
def side_effect_function(x):
  y = x ** 2
  results.append(y) # Side effect: modifying a global variable
  return y




# Call it
z = side_effect_function(3.)




# What happens if we try to use the result?
# print(results[0] + 1)
# This would raise an UnexpectedTracerError!

The side effect (results.append(y)) happens at trace time. This means a tracer object representing y gets put into the results list. When you later try to access results[0] outside the jit-ted function, you're pulling a live tracer out of the computation graph, which JAX prohibits.

The fix is always to make the function pure. If you need to get a value out, return it explicitly.

4. A Gallery of Other Gotchas

With the core principles of immutability, static tracing, and purity in mind, we can quickly diagnose a few other common pitfalls.

• Out-of-Bounds Indexing: In NumPy, indexing an array out of bounds raises an IndexError. In JAX, raising errors from compiled code on an accelerator is difficult. Instead, JAX has defined behavior: out-of-bounds writes are dropped, and reads are "clamped" to the array's boundary (e.g., reading from index 100 in a 10-element array will return the element at index 9). This can lead to silent errors if you're not careful.
• Using numpy instead of jax.numpy: Inside a jit-ted function, you must use jnp functions. Standard np functions don't know how to handle JAX tracers and will fail.
• Casting Tracers to Python Types: Calling int(x), float(x), or x.item() on a traced value x will trigger a ConcretizationTypeError. The operation demands a concrete value, but it only has a tracer. You must keep values as JAX arrays within the transformed function.

Conclusion

Navigating the transition from NumPy to JAX involves internalizing a few new, powerful constraints. These aren't arbitrary rules but are essential for enabling JAX's performance and features.

Let's summarize the key takeaways:

• Immutability is Law: JAX arrays cannot be changed in-place. Use the functional .at[...].set(...) pattern to create updated copies.
• Think Like the Tracer: Functions to be transformed must have a static computation graph. Any Python control flow (if, for) that depends on the values of your data will fail during tracing. Use JAX-native control flow like jnp.where instead.
• Embrace Purity: All data a function needs must be passed in as arguments. All results it produces must be explicitly returned. Side effects, from modifying globals to printing, will happen at trace time and lead to errors or confusing behavior.

This mental model—of a static graph built from pure functions operating on immutable data—is the key to writing effective JAX code.

In our next lesson, we will focus exclusively on the functional update API, mastering the use of .at[].set() and its variants to handle array modifications in a clean and JAX-idiomatic way.