JAX vs. NumPy vs. PyTorch: Execution, Immutability, and Purity

1. A New Paradigm: Functional vs. Object-Oriented

PyTorch, like many Python libraries, embraces an object-oriented, imperative style. You create model objects that encapsulate both behavior (forward method) and state (parameters in self.layer). You then execute operations step-by-step, and these operations can mutate the state of your objects, for example, when an optimizer updates model weights in place.

JAX is built on a different foundation: functional programming. This paradigm treats computation as the evaluation of mathematical functions and avoids changing state and mutable data. The benefits of this strictness are significant, enabling a compiler to perform powerful optimizations and making parallelism more straightforward.

To start, let's get a high-level overview from a guide written specifically for PyTorch developers.

As the article mentions, the core of JAX's functional approach is the concept of pure functions. A function is considered pure if it satisfies two main conditions:

1. Its output depends only on its explicit inputs.
2. It has no "side effects," meaning it doesn't modify any external state, print to the console, or read from a global variable.

This is a stark contrast to a typical Python function or method that might modify an object's attribute (self.value = ...) or a global variable.

The following video from Google for Developers provides a concise and clear explanation of what functional purity means in the context of JAX.

2. Immutability and Explicit State

The principle of functional purity leads directly to JAX's second key characteristic: immutability. Since a pure function cannot modify its inputs (as that would be a side effect), JAX arrays are immutable. You cannot change their values in-place.

In PyTorch or NumPy, you would typically write:

# PyTorch/NumPy style
import numpy as np
a = np.arange(5)
a[1] = 99 # In-place mutation
print(a) # >> [ 0 99  2  3  4]

Attempting this in JAX will result in an error. Instead, JAX provides a specific functional syntax, .at[...].set(), which returns a new array with the updated value, leaving the original array untouched.

# JAX style
import jax.numpy as jnp
a = jnp.arange(5)
b = a.at[1].set(99) # Returns a new, updated array
print(a) # >> [0 1 2 3 4] (original is unchanged)
print(b) # >> [ 0 99  2  3  4]

While this might seem less efficient at first glance, the XLA compiler often optimizes these operations "under the hood," fusing them to avoid unnecessary memory allocation.

This concept extends beyond simple arrays to all aspects of state management. In PyTorch, a model nn.Module holds its own parameters. In JAX, because of purity, state (like model weights, optimizer state, or even random number generator state) cannot be stored inside an object. It must be handled explicitly: passed as an argument to a function and returned as an output if it has been updated.

Let's delve deeper into these practical consequences with a more detailed article.

3. The Execution Model: Tracing, jaxpr, and XLA

So, why enforce these strict rules of purity and immutability? The payoff is in the execution model. JAX is not just an interpreter that runs your Python code line by line. Instead, it uses a process called tracing.

When you apply a transformation like jax.jit (Just-In-Time compilation) to a function, JAX doesn't run the function immediately. It calls it once with abstract "tracer" objects that record every operation performed. The sequence of operations is captured in a simple, intermediate representation called a jaxpr (JAX expression).

This jaxpr is a pure, static computation graph, disentangled from the complexities of Python. Because it's guaranteed to be free of side effects, JAX can safely hand it over to the XLA (Accelerated Linear Algebra) compiler. XLA performs powerful, domain-specific optimizations—like fusing multiple operations into a single kernel, optimizing memory layouts, and compiling the graph into highly efficient machine code for your specific hardware (CPU, GPU, or TPU).

The following video explains this tracing mechanism beautifully.

This trace-then-compile execution model is the fundamental difference between JAX and PyTorch's default eager execution. While PyTorch 2.0's torch.compile() introduces a similar compilation path, JAX was designed around this principle from the very beginning.

The diagram below shows the entire AI execution stack for both JAX and PyTorch, illustrating how your high-level model code is progressively lowered to machine code.

In this diagram, you can clearly see the JAX path: your code is lowered into a JAXPR, which is then converted to an XLA-compatible format (HLO/StableHLO). The XLA compiler takes it from there to generate optimized code for the target device. This clean, direct path to a powerful compiler is the primary reason for JAX's impressive performance.

Summary

This lesson covered the conceptual bedrock of JAX. The differences between JAX and PyTorch are not just superficial API changes; they represent a fundamental split in programming paradigms.

• Functional Purity: JAX is built around pure functions, which have no side effects and whose output depends only on their inputs. This is unlike PyTorch's object-oriented model where objects encapsulate and mutate state.
• Immutability: A direct consequence of purity is that JAX arrays are immutable. Operations that would modify an array in-place in NumPy/PyTorch instead return a new, modified array in JAX (e.g., .at[].set()).
• Execution Model: JAX uses a trace-and-compile model. It traces your Python code to generate a jaxpr, a simple intermediate representation. This jaxpr is then optimized and compiled to fast machine code by the XLA compiler. This contrasts with PyTorch's default eager execution model.

These three pillars are interconnected. Functional purity enables reliable tracing, which in turn produces a clean computation graph for the XLA compiler to optimize, delivering the high performance JAX is known for.

In our next lesson, we will move from theory to practice. You will get your hands dirty with the jax.numpy API, creating and manipulating JAX arrays and seeing firsthand how these principles of immutability and functional design play out in code.