Pure Functions: The Cornerstone of JAX

Introduction

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.

What Are Pure Functions?

f(x) = x^2

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Issues of Side Effects: Global State

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.

Python

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)")

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:

text

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

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.

Building Pure Alternatives

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:

Python

# 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)")

# 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:

text

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

Python

# 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}")

# 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:

text

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

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.

Pure Functions with JAX Arrays

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:

Python

# 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}")

# 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:

text

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

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:

Python

# 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)}")

# 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:

text

Is output1 identical to output2? True

Is output1 identical to output2? True

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

Introduction

What Are Pure Functions?

Join the 1M+ learners on CodeSignal

Issues of Side Effects: Global State

Issues of Side Effects: In-Place Modification Mutates Inputs

Building Pure Alternatives

Pure Functions with JAX Arrays

Why JAX Demands Purity

Conclusion and Next Steps