Pure Functions: The Cornerstone of JAX

Pure Functions: The Cornerstone of JAX | CodeSignal Learn

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

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:

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.
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 $9$ , regardless of when or how many times we call it. The function doesn't modify $x$ or affect anything else in the mathematical universe — it simply computes and returns a result.

This purity might seem restrictive at first, but it's what enables JAX to perform its magic. When JAX knows that a function is pure, it can safely optimize it, parallelize it, and even differentiate it automatically. These transformations would be impossible or unreliable if functions could have unpredictable side effects.

In the context of numerical computing and machine learning, pure functions make our code more predictable, testable, and optimizable. They eliminate many sources of bugs that arise from unexpected state changes and make it easier to reason about our programs' behavior.

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. Pythonimport 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: textInitial 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.

Issues of Side Effects: In-Place Modification Mutates Inputs

Another common form of impurity involves modifying input arguments in place. Let's see how this can cause problems: Python# 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: textOriginal 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.

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)") The output shows predictable behavior without side effects: textcurrent_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}") This pure version preserves our original data: textOriginal 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}") The output demonstrates how JAX arrays naturally support pure function design: textOriginal 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)}") As expected, the function produces identical results: textIs output1 identical to output2? True This deterministic behavior is exactly what JAX needs to perform its transformations safely and efficiently.