JAX Arrays: The Immutable Successor to NumPy

JAX Arrays: The Immutable Successor to NumPy | CodeSignal Learn

Introduction

What is JAX and Why Should You Care?

Before starting, you might be wondering: why should I care about JAX? JAX stands as one of the most exciting developments in numerical computing and machine learning. Developed by Google Research, this framework brings together the familiar NumPy API we know and love with cutting-edge performance optimizations and functional programming principles that make it incredibly powerful for scientific computing and deep learning. At its core, JAX is built on three fundamental pillars that set it apart from traditional NumPy: First, JAX embraces immutability — once we create an array, we cannot modify it in place. While this might seem restrictive at first, immutability enables powerful optimizations and makes our code more predictable and easier to reason about, especially in parallel computing environments. Second, JAX provides automatic differentiation (autodiff) that can compute gradients of any function we write. This capability is essential for machine learning, where we need to compute gradients to train our models. Unlike manual differentiation or finite differences, JAX's autodiff is both fast and numerically stable. Third, JAX includes advanced compilation capabilities that can dramatically accelerate our computations. We'll explore these performance features throughout our course, building from simple concepts to powerful optimizations. What makes JAX particularly compelling is that it maintains NumPy's familiar interface while adding these superpowers. If you can write NumPy code, you're already 80% of the way to writing JAX code! The learning curve is gentle, but the performance gains and capabilities you'll unlock are transformational.

Join the 1M+ learners on CodeSignal

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

Creating JAX Arrays

Let's start our hands-on journey by learning how to create JAX arrays. Just as NumPy serves as the foundation for scientific computing in Python, JAX arrays serve as the foundation for everything we'll build in this course. JAX provides the jax.numpy module (commonly imported as jnp) that mirrors NumPy's functionality. We can create JAX arrays from Python lists, NumPy arrays, or by using any of the familiar array creation functions: Pythonimport jax.numpy as jnp import numpy as np # Creating JAX arrays from Python lists py_list = [1.0, 2.0, 3.0] jax_array_from_list = jnp.array(py_list) print(f"JAX array from list: {jax_array_from_list}") # Creating JAX arrays from NumPy arrays numpy_array = np.array([4.0, 5.0, 6.0]) jax_array_from_numpy = jnp.array(numpy_array) print(f"JAX array from NumPy array: {jax_array_from_numpy}") When we run this code, we'll see that JAX arrays look and behave very similarly to NumPy arrays in terms of their string representation. The key difference lies under the hood — JAX arrays are designed for functional programming and come with the immutability guarantee we'll explore shortly. Notice that we're using jnp.array() instead of np.array(). This is our gateway into the JAX ecosystem. The function signature and behavior are nearly identical to NumPy's array() function, making the transition seamless for anyone familiar with NumPy.

Basic Operations with JAX Arrays

Now that we can create JAX arrays, let's explore how to perform computations with them. JAX maintains the same intuitive syntax as NumPy for mathematical operations, making our transition smooth and natural. Python# Element-wise operations arr1 = jnp.array([1.0, 2.0, 3.0]) arr2 = jnp.array([4.0, 5.0, 6.0]) element_wise_sum = arr1 + arr2 # Broadcasting and addition work as expected print(f"Element-wise sum: {element_wise_sum}") # Matrix operations mat1 = jnp.array([[1.0, 2.0], [3.0, 4.0]]) mat2 = jnp.array([[5.0, 6.0], [7.0, 8.0]]) matrix_product = jnp.dot(mat1, mat2) # Standard matrix multiplication print(f"Matrix product:\n{matrix_product}") When we execute this code, we'll see output like: Element-wise sum: [5. 7. 9.] Matrix product: [[19. 22.] [43. 50.]] The beauty of JAX lies in this familiar interface. If you've written NumPy code before, this feels completely natural. We can perform element-wise addition with the + operator and matrix multiplication using jnp.dot(), just as we would in NumPy. Behind the scenes, however, JAX is preparing these operations for potential optimization and making them compatible with automatic differentiation. This means that any function we write using these basic operations can later be automatically differentiated or optimized for enhanced performance.

The Power of Immutability

Here's where JAX diverges significantly from NumPy, and this difference is crucial to understand. JAX arrays are immutable, meaning that once created, they cannot be modified. This design choice might feel unfamiliar at first, but it's fundamental to JAX's functional programming paradigm and enables many of its powerful features. Let's see what happens when we try to modify a JAX array in the traditional NumPy way: Python# Demonstrating immutability immutable_jax_array = jnp.array([10.0, 20.0, 30.0]) print(f"Original JAX array: {immutable_jax_array}") try: # This will raise an error - JAX arrays cannot be modified in place immutable_jax_array[0] = 100.0 except TypeError as e: print(f"Error when trying to modify JAX array in-place: {e}") When we run this code, JAX will raise a helpful error message telling us that arrays are immutable and pointing us toward the correct approach: Original JAX array: [10. 20. 30.] Error when trying to modify JAX array in-place: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method As suggested by the error message, instead of modifying arrays in place, JAX provides the .at[] method for functional updates: Python# The correct way to update JAX arrays - functional style updated_jax_array = immutable_jax_array.at[0].set(100.0) print(f"Original JAX array (unchanged): {immutable_jax_array}") print(f"New JAX array after update: {updated_jax_array}") This approach creates a new array with our desired changes while leaving the original array completely unchanged. The output demonstrates this clearly: Original JAX array (unchanged): [10. 20. 30.] New JAX array after update: [100. 20. 30.]

Immutable JAX against Mutable NumPy

To contrast this with NumPy's mutable behavior, let's recall how traditional NumPy arrays behave: Python# NumPy arrays are mutable for comparison mutable_numpy_array = np.array([10.0, 20.0, 30.0]) print(f"Original NumPy array: {mutable_numpy_array}") mutable_numpy_array[0] = 100.0 # Direct in-place modification print(f"NumPy array after in-place modification: {mutable_numpy_array}") The NumPy array changes in place: Original NumPy array: [10. 20. 30.] NumPy array after in-place modification: [100. 20. 30.] This mutable behavior can lead to unexpected side effects in complex programs. JAX's immutability eliminates these surprises and makes our code more predictable and safer for parallel execution.

Performance Insights: A Glimpse of JAX's Potential

One of the most notable performance optimization of JAX is known as just-in-time (JIT) compilation. While we'll explore JAX's JIT mechanism in much greater detail later in our course, let's get a taste of what's possible with a simple benchmark. This will help us understand why JAX has become so popular in the scientific computing community. Pythonimport time # Simple benchmark comparing NumPy and JAX size = 10000000 np_arr = np.random.rand(size).astype(np.float32) # Large array for meaningful timing jax_arr = jnp.array(np_arr) # Convert to JAX array # NumPy timing start_time = time.time() np_sum = np.sum(np_arr) _ = np_sum.item() # Ensure computation is complete end_time = time.time() print(f"NumPy sum time: {end_time - start_time:.6f} seconds") # JAX timing with JIT jax_sum_op = jax.jit(jnp.sum) # First run includes compilation start_time = time.time() jax_sum = jax_sum_op(jax_arr) jax_sum.block_until_ready() end_time = time.time() print(f"JAX (JIT) sum time (first run, includes compilation): {end_time - start_time:.6f} seconds") # Second run uses cached compilation start_time = time.time() jax_sum = jax_sum_op(jax_arr) jax_sum.block_until_ready() end_time = time.time() print(f"JAX (JIT) sum time (second run): {end_time - start_time:.6f} seconds") This benchmark reveals something fascinating about JAX's execution model and performance characteristics: NumPy sum time: 0.003732 seconds JAX (JIT) sum time (first run, includes compilation): 0.100887 seconds JAX (JIT) sum time (second run): 0.001683 seconds Notice the interesting pattern here! The first JAX run is significantly slower because it includes compilation overhead — JAX is analyzing our function and optimizing it for the specific hardware. However, the second run is actually faster than NumPy, demonstrating the power of JIT compilation. This pattern is crucial to understand: JAX trades some initial setup time for dramatically improved performance on subsequent runs. For machine learning workloads where we repeatedly call the same functions (like forward passes through a neural network), this trade-off is incredibly beneficial. The .block_until_ready() call ensures our timing is accurate because JAX operations are asynchronous by default. This asynchronous execution is part of what enables JAX's performance optimizations.

JAX Arrays: The Immutable Successor to NumPy

Introduction

What is JAX and Why Should You Care?

Join the 1M+ learners on CodeSignal

Creating JAX Arrays

Basic Operations with JAX Arrays

The Power of Immutability

Immutable JAX against Mutable NumPy

Performance Insights: A Glimpse of JAX's Potential

Conclusion and Next Steps