JAX Arrays: The Immutable Successor to NumPy

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

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

# 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

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

# 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.]

Original JAX array (unchanged): [10. 20. 30.]
New JAX array after update: [100.  20.  30.]

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.

Python

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

import 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

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.

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