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.

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.

import jax.numpy as jnp

# Impure function: modifies global state
global_counter = 0
def impure_increment_global():
    global global_counter
    global_counter += 1  # Side effect: modifying global variable
    return global_counter

# Let's see the side effects in action
print(f"Initial global_counter: {global_counter}")
result1 = impure_increment_global()
print(f"Result 1: {result1}, global_counter: {global_counter}")
result2 = impure_increment_global()
print(f"Result 2: {result2}, global_counter: {global_counter} (Side-effect observed)")

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

Initial global_counter: 0
Result 1: 1, global_counter: 1
Result 2: 2, global_counter: 2 (Side-effect observed)

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:

# Impure function: modifies input list in-place
def impure_scale_list_inplace(data_list, scalar):
    for i in range(len(data_list)):
        data_list[i] *= scalar  # Side effect: modifying input argument
    return data_list  # Returns the same modified object

# Demonstrating the side effect
my_list = [1, 2, 3]
print(f"Original list: {my_list}")
modified_list = impure_scale_list_inplace(my_list, 2)
print(f"List after impure_scale_list_inplace: {my_list} (Input modified)")
print(f"Returned list is same object: {modified_list is my_list}")

The output reveals the problematic side effect:

Original list: [1, 2, 3]
List after impure_scale_list_inplace: [2, 4, 6] (Input modified)
Returned list is same object: True

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:

# Pure version: no global state dependency
def pure_increment(value):
    return value + 1  # Simply returns new value, no side effects

# Using the pure function
current_value = 0
current_value = pure_increment(current_value)
print(f"current_value after first pure_increment: {current_value}")
current_value = pure_increment(current_value)
print(f"current_value after second pure_increment: {current_value} (No side-effects)")

The output shows predictable behavior without side effects:

current_value after first pure_increment: 1
current_value after second pure_increment: 2 (No 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:

# Pure version for list: creates new list instead of modifying input
def pure_scale_list(data_list, scalar):
    return [x * scalar for x in data_list]  # List comprehension creates new list

# Demonstrating pure behavior
my_list_for_pure = [1, 2, 3]
print(f"Original list: {my_list_for_pure}")
scaled_list = pure_scale_list(my_list_for_pure, 2)
print(f"Original list unchanged: {my_list_for_pure}")
print(f"New scaled list: {scaled_list}")

This pure version preserves our original data:

Original list: [1, 2, 3]
Original list unchanged: [1, 2, 3]
New scaled list: [2, 4, 6]

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:

# Pure function for JAX arrays
def pure_scale_jax_array(data_array, scalar):
    return data_array * scalar  # Creates new array due to immutability

# Working with JAX arrays
jax_arr = jnp.array([10.0, 20.0, 30.0])
scalar_val = 3.0
print(f"Original JAX array: {jax_arr}")

scaled_jax_arr = pure_scale_jax_array(jax_arr, scalar_val)
print(f"Scaled JAX array: {scaled_jax_arr}")
print(f"Original JAX array unchanged: {jax_arr}")

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

Original JAX array: [10. 20. 30.]
Scaled JAX array: [30. 60. 90.]
Original JAX array unchanged: [10. 20. 30.]

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:

# Demonstrating determinism of pure functions
output1 = pure_scale_jax_array(jax_arr, scalar_val)
output2 = pure_scale_jax_array(jax_arr, scalar_val)
print(f"Is output1 identical to output2? {jnp.array_equal(output1, output2)}")

As expected, the function produces identical results:

Is output1 identical to output2? True

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.