Welcome back to "JAX Fundamentals: NumPy Power-Up"! We're making excellent progress on our journey to mastering JAX. In our previous lesson, we explored JAX arrays and discovered how their immutability sets them apart from traditional NumPy arrays. We learned that JAX arrays cannot be modified in place and instead require functional updates using the .at[] method.

Today, we're building upon that foundation to explore another fundamental concept that makes JAX so powerful: pure functions. As you may recall from our previous lesson, JAX embraces functional programming principles, and immutability was our first taste of this paradigm. Pure functions represent the next crucial step in understanding why JAX is designed the way it is.

This lesson will help us understand what pure functions are, why they matter, and how they enable JAX's most powerful features, such as automatic differentiation and just-in-time compilation. By the end of this lesson, we'll be able to identify impure functions, refactor them into pure alternatives, and understand why JAX's transformations rely so heavily on function purity.

Before we dive into code examples, let's establish a clear understanding of what makes a function pure. In functional programming, a pure function is one that satisfies two essential criteria:

1. Deterministic behavior: Given the same inputs, a pure function will always produce the same outputs. There's no randomness, no dependency on external state, and no variation between calls.

2. No side effects: A pure function doesn't modify anything outside of itself. It doesn't change global variables, modify its input arguments, write to files, print to the console, or interact with any external systems.

Think of pure functions like mathematical functions. When we write $f(x) = x^2$, we expect that $f(3)$ will always equal , regardless of when or how many times we call it. The function doesn't modify or affect anything else in the mathematical universe — it simply computes and returns a result.

To understand the importance of pure functions, let's first examine what happens when functions are impure. Impure functions create side effects that can make our code unpredictable and difficult to optimize. Let's start with a common example: a function that modifies global state.

When we run this code, we can observe how the function's behavior depends on external state:

Notice how impure_increment_global() returns different values each time we call it, even though we're not passing any arguments! This violates the deterministic behavior requirement of pure functions. The function's output depends on the current value of global_counter, and calling the function modifies this global state.

Another common form of impurity involves modifying input arguments in place. Let's see how this can cause problems:

The output reveals the problematic side effect:

The function has modified our original list! This can lead to subtle bugs in larger programs where we might not expect our data to change after passing it to a function.

Now let's refactor these impure functions into pure alternatives that eliminate side effects while maintaining the same functionality. Pure functions solve the problems we just observed by creating new values instead of modifying existing state.

For our global counter example, we can create a pure version that takes the current value as an input parameter:

The output shows predictable behavior without side effects:

Notice how pure_increment(1) will always return 2, regardless of when we call it or what happened before. The function is completely deterministic and doesn't modify any external state.

For our list scaling example, we can create a pure version that returns a new list:

This pure version preserves our original data:

The key insight here is that instead of modifying the input, we create and return a new list with the scaled values. Our original data remains untouched, eliminating the side effect that could cause problems elsewhere in our program.

Now let's see how these pure function principles apply specifically to JAX arrays. Since JAX arrays are immutable (as we learned in our previous lesson), they naturally encourage pure function design. Let's create a pure function that works with JAX arrays:

The output demonstrates how JAX arrays naturally support pure function design:

Because JAX arrays are immutable, the multiplication operation data_array * scalar automatically creates a new array rather than modifying the input. This makes it nearly impossible to accidentally create impure functions when working with JAX arrays.

We can also demonstrate the deterministic nature of pure functions:

As expected, the function produces identical results:

This deterministic behavior is exactly what JAX needs to perform its transformations safely and efficiently.

You might wonder: why does JAX care so much about pure functions? The answer lies in the powerful transformations that JAX provides, which we'll explore in detail in upcoming lessons. These transformations — like automatic differentiation and just-in-time compilation — rely fundamentally on the predictable behavior that pure functions guarantee.

When JAX compiles a function for optimization, it needs to know that the function will behave exactly the same way every time it's called with the same inputs: if a function could modify global state or behave differently based on external factors, JAX's optimizations could produce incorrect results or fail entirely. Similarly, automatic differentiation works by analyzing the mathematical operations within a function. If a function has side effects — like modifying global variables or printing output — these side effects don't have mathematical derivatives, and including them in the differentiation process would be meaningless or problematic.

Pure functions also enable safe parallelization: when JAX distributes computations across multiple processors or devices, it needs to know that functions won't interfere with each other through shared state modifications. Pure functions guarantee this independence.

Think of purity as the foundation that makes JAX's magic possible. Without pure functions, the automatic differentiation, compilation, and parallelization features that make JAX so powerful would be unreliable or impossible to implement safely.

Congratulations! We've explored the fundamental concept of pure functions and discovered why they're absolutely essential for JAX's operation. We've learned to identify impure functions by looking for side effects like global state modification and in-place argument changes, and we've practiced refactoring these into pure alternatives that create new values instead of modifying existing ones. The key insights are that pure functions are deterministic and have no side effects, JAX's immutable arrays naturally encourage pure function design, and pure functions enable JAX's powerful transformations.

As we continue our journey through JAX, we'll see how pure functions become the building blocks for more advanced concepts. In our upcoming lessons, we'll explore automatic differentiation and just-in-time compilation — two transformations that absolutely depend on the function purity we've mastered today. The discipline of writing pure functions might feel constraining at first, but it's this constraint that unlocks JAX's extraordinary capabilities.