JAX Arrays: NumPy & PyTorch Compared

1. The jax.numpy API: A Familiar Entry Point

JAX is designed to be approachable, especially for those with a background in the scientific Python stack. Its primary array manipulation interface, jax.numpy, is intentionally modeled after the NumPy API. For a vast number of operations, you can think of it as a near drop-in replacement that enables execution on accelerators and unlocks JAX's transformations.

Let's begin with a quick look at how closely jax.numpy mirrors its namesake.

As you can see, the learning curve for basic array creation and arithmetic is virtually flat. You can leverage your existing knowledge of NumPy's API for a large portion of your work in JAX. However, beneath this familiar surface lie crucial distinctions that we will now explore in detail.

2. Creating and Inspecting jax.Array Objects

The jax.Array is the JAX counterpart to torch.Tensor. Let's review the common ways to create these arrays and examine their core attributes, noting the important differences from PyTorch tensors.

The reading highlights three critical differences between jax.Array and torch.Tensor:

1. JAX arrays are immutable.
2. The default dtypes are 32-bit.
3. The default device is the first available accelerator.

Let's dissect each of these, as they have significant practical implications.

Immutability in Practice

In the previous lesson, we discussed immutability as a consequence of functional purity. In PyTorch or NumPy, you can modify an array's contents directly. In JAX, this is disallowed and will raise a TypeError.

import jax.numpy as jnp
import numpy as np




# NumPy: Mutable
a_np = np.arange(5.)
a_np[1] = 99.
print(f"NumPy array (mutated): {a_np}")




# JAX: Immutable
a_jnp = jnp.arange(5.)
try:
    a_jnp[1] = 99.
except TypeError as e:
    print(f"\nJAX raises TypeError: {e}")




# The functional way in JAX
b_jnp = a_jnp.at[1].set(99.)
print(f"Original JAX array (unchanged): {a_jnp}")
print(f"New JAX array (updated): {b_jnp}")

This might seem inefficient, as if a full copy of the array is being made. However, this is where the XLA compiler comes in. The following video explains how JAX handles these "updates" and how the compiler can optimize them under the hood.

Default Data Types

PyTorch defaults to 32-bit floating-point numbers (torch.float32), which aligns well with GPU hardware. NumPy, on the other hand, is often more aggressive in promoting to 64-bit precision (numpy.float64). JAX follows the PyTorch convention and defaults to 32-bit precision. This is a deliberate design choice for performance on accelerators, which are highly optimized for 32-bit (and 16-bit) arithmetic.

You can, of course, specify dtype=jnp.float64 if you need higher precision, but it's important to be aware of the default. You can enable 64-bit precision globally in JAX, but it's generally recommended to be explicit about dtype when higher precision is required.

Default Device Placement

This is perhaps the most significant daily workflow difference from PyTorch. In PyTorch, you must explicitly move tensors to a device: my_tensor.to('cuda'). If you forget, you get a RuntimeError when you try to operate on tensors residing on different devices.

JAX takes a different approach. By default, arrays are created on the first available accelerator (e.g., 'gpu:0'). An array without an explicit device assignment is considered uncommitted; it resides on the default device but can be implicitly moved to another device if an operation requires it. This simplifies code by removing the need for frequent .to(device) calls.

This video provides an excellent demonstration of JAX's accelerator-agnostic design and the performance implications of data placement.

To explicitly control device placement, you use jax.device_put():

import jax




# List available devices
print(f"Available JAX devices: {jax.devices()}")




# Create a CPU device object
cpu = jax.devices('cpu')[0]




# Create an array and explicitly place it on the CPU
x_cpu = jax.device_put(jnp.arange(5), device=cpu)

print(f"\nArray device: {x_cpu.device()}")

jax.device_put is the JAX equivalent of PyTorch's tensor.to(device).

3. A Visual Contrast: Imperative vs. Functional

The differences in array types are a direct reflection of the underlying design philosophies of the frameworks. PyTorch builds a computation graph dynamically as you execute operations on its mutable tensor objects. JAX, in contrast, creates new functions that compute results from immutable inputs.

While we will cover automatic differentiation in detail in a future module, the following diagram provides a powerful visual summary of this core difference in how PyTorch and JAX operate on their respective array objects to compute gradients.

This image neatly encapsulates the paradigm shift: PyTorch operations mutate a graph attached to the tensor objects, while JAX operations are composed within pure functions that are then transformed.

Key Takeaways

This lesson provided a practical introduction to the jax.Array, always in the context of your existing expertise with PyTorch and NumPy.

• Familiar API, Different Semantics: jax.numpy provides a familiar API, but the underlying jax.Array objects behave differently from NumPy arrays and PyTorch tensors.
• Immutability is Law: You cannot change a JAX array in-place. Use the .at[...].set() syntax to create a new, updated array. This is a cornerstone of JAX's functional design, enabling robust compiler optimizations.
• Sensible Defaults for Performance: JAX defaults to 32-bit precision and places arrays on the primary accelerator by default, streamlining code for high-performance computing.
• Explicit is Better than Implicit (for Control): While default placement is convenient, jax.device_put() gives you explicit control when needed, analogous to PyTorch's .to(device).

In our next lesson, we will continue exploring device management. We'll look more closely at how to use jax.device_put and inspect array properties to ensure your data is where you expect it to be, setting the stage for more complex, single-accelerator programs.