Chapter 2: NumPy-like API & Arrays

Overview

JAX arrays behave like NumPy arrays with key differences. This chapter covers JAX's array API and common operations.

---

JAX Arrays vs NumPy Arrays

Similarities

import jax.numpy as jnp
import numpy as np

# Both have similar APIs
x_jax = jnp.array([1, 2, 3])
x_np = np.array([1, 2, 3])

# Both support operations
print(jnp.sum(x_jax))        # 6
print(np.sum(x_np))          # 6

print(jnp.sin(x_jax))        # [0.84... 0.90... 0.14...]
print(np.sin(x_np))          # [0.84... 0.90... 0.14...]

# Both support shapes
print(x_jax.shape)           # (3,)
print(x_np.shape)            # (3,)

Key Differences

Feature	NumPy	JAX
Mutability	Mutable (modify in-place)	Immutable
Device	CPU only	CPU/GPU/TPU
Dtype default	float64	float32
Indexing updates	x[0] = 5	x.at[0].set(5)
Tracing	None	Traces with symbolic values

---

Creating JAX Arrays

From Lists and Tuples

import jax.numpy as jnp

# From list
x = jnp.array([1, 2, 3])
print(x)  # [1 2 3]

# From nested list (2D array)
A = jnp.array([[1, 2], [3, 4]])
print(A.shape)  # (2, 2)

# From tuple
x = jnp.array((1, 2, 3))
print(x)  # [1 2 3]

# Specify dtype
x = jnp.array([1, 2, 3], dtype=jnp.float32)
print(x.dtype)  # float32

Initialization Functions

import jax.numpy as jnp

# Zeros
z = jnp.zeros((3, 4))           # 3x4 matrix of zeros
z = jnp.zeros(5, dtype=jnp.int32)  # 5-element int array

# Ones
o = jnp.ones((2, 3))            # 2x3 matrix of ones

# Identity
I = jnp.eye(3)                  # 3x3 identity matrix

# Range
r = jnp.arange(10)              # [0, 1, 2, ..., 9]
r = jnp.arange(0, 10, 2)        # [0, 2, 4, 6, 8]

# Linspace
l = jnp.linspace(0, 1, 5)       # [0., 0.25, 0.5, 0.75, 1.]

# Constants
c = jnp.full((3,), 5.0)         # [5., 5., 5.]

Random Arrays

import jax
import jax.numpy as jnp

# JAX has special random number handling (explained later)
key = jax.random.PRNGKey(0)     # Create random key

# Normal distribution
x = jax.random.normal(key, (5,))           # 5 random normal values
x = jax.random.normal(key, (3, 4))         # 3x4 normal matrix

# Uniform distribution [0, 1)
x = jax.random.uniform(key, (5,))

# Custom range
x = jax.random.uniform(key, (5,), minval=0, maxval=10)

# Integer
x = jax.random.randint(key, (5,), minval=0, maxval=100)

---

Array Properties

Shape and Size

import jax.numpy as jnp

A = jnp.array([[1, 2, 3], [4, 5, 6]])

# Shape
print(A.shape)      # (2, 3)
print(A.ndim)       # 2 (number of dimensions)
print(A.size)       # 6 (total elements)

# Access dimensions
rows, cols = A.shape
print(f"{rows} rows, {cols} columns")  # 2 rows, 3 columns

Data Types

import jax.numpy as jnp

# Check dtype
x = jnp.array([1.0, 2.0])
print(x.dtype)      # float32 (JAX default)

# Convert dtype
x_int = x.astype(jnp.int32)
x_float64 = x.astype(jnp.float64)

# Common dtypes in JAX
int_array = jnp.array([1, 2, 3], dtype=jnp.int32)
float_array = jnp.array([1, 2, 3], dtype=jnp.float32)
bool_array = jnp.array([True, False], dtype=jnp.bool_)

# Note: JAX defaults to float32 (not float64 like NumPy!)
x = jnp.array([1.0])
print(x.dtype)  # float32 (not float64!)

Why float32 default? For GPU/TPU efficiency and memory usage.

---

Element-Wise Operations

Arithmetic

import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
y = jnp.array([4.0, 5.0, 6.0])

# Element-wise addition
z = x + y              # [5. 7. 9.]

# Element-wise subtraction
z = x - y              # [-3. -3. -3.]

# Element-wise multiplication
z = x * y              # [4. 10. 18.]

# Element-wise division
z = y / x              # [4. 2.5 2.]

# Power
z = x ** 2             # [1. 4. 9.]

# Scalar operations work too
z = x * 2              # [2. 4. 6.]
z = x + 10             # [11. 12. 13.]

Trigonometric Functions

import jax.numpy as jnp

x = jnp.array([0., jnp.pi/2, jnp.pi])

print(jnp.sin(x))      # [0. 1. 0.]
print(jnp.cos(x))      # [1. 0. -1.]
print(jnp.tan(x))      # [0. inf -0.]

# Inverse
x = jnp.array([0.5, 0.0, -0.5])
print(jnp.arcsin(x))   # [-π/6, 0, -π/6]
print(jnp.arccos(x))   # [π/3, π/2, 2π/3]

Exponentials and Logarithms

import jax.numpy as jnp

x = jnp.array([1., 2., 3.])

print(jnp.exp(x))      # [2.71... 7.38... 20.08...]
print(jnp.log(x))      # [0. 0.69... 1.09...]
print(jnp.log10(x))    # [0. 0.30... 0.47...]
print(jnp.sqrt(x))     # [1. 1.41... 1.73...]

Rounding

import jax.numpy as jnp

x = jnp.array([1.23, 4.56, 7.89])

print(jnp.round(x))    # [1. 5. 8.]
print(jnp.floor(x))    # [1. 4. 7.]
print(jnp.ceil(x))     # [2. 5. 8.]

Absolute Value and Sign

import jax.numpy as jnp

x = jnp.array([-3., -1., 0., 1., 3.])

print(jnp.abs(x))      # [3. 1. 0. 1. 3.]
print(jnp.sign(x))     # [-1. -1. 0. 1. 1.]

---

Reduction Operations

Reductions combine elements along axes.

Sum and Mean

import jax.numpy as jnp

A = jnp.array([[1., 2., 3.],
               [4., 5., 6.]])

# Sum all elements
print(jnp.sum(A))           # 21.

# Sum along axis 0 (rows → columns)
print(jnp.sum(A, axis=0))   # [5. 7. 9.]

# Sum along axis 1 (columns → rows)
print(jnp.sum(A, axis=1))   # [6. 15.]

# Mean
print(jnp.mean(A))          # 3.5
print(jnp.mean(A, axis=0))  # [2.5 3.5 4.5]

Min and Max

import jax.numpy as jnp

x = jnp.array([3., 1., 4., 1., 5.])

print(jnp.min(x))           # 1.
print(jnp.max(x))           # 5.
print(jnp.argmin(x))        # 1 (index of minimum)
print(jnp.argmax(x))        # 4 (index of maximum)

# 2D example
A = jnp.array([[1., 2., 3.],
               [4., 5., 6.]])

print(jnp.min(A, axis=1))   # [1. 4.]
print(jnp.max(A, axis=1))   # [3. 6.]

Std and Var

import jax.numpy as jnp

x = jnp.array([1., 2., 3., 4., 5.])

print(jnp.std(x))           # ~1.414 (standard deviation)
print(jnp.var(x))           # 2.0 (variance)

---

Reshaping and Transposing

Reshape

import jax.numpy as jnp

x = jnp.arange(12)  # [0, 1, 2, ..., 11]
print(x.shape)      # (12,)

# Reshape to 2D
A = x.reshape((3, 4))
print(A.shape)      # (3, 4)
print(A)            # [[0 1 2 3]
                    #  [4 5 6 7]
                    #  [8 9 10 11]]

# Reshape to 3D
B = x.reshape((2, 3, 2))
print(B.shape)      # (2, 3, 2)

# Flatten
flat = A.flatten()
print(flat.shape)   # (12,)
print(flat)         # [0 1 2 ... 11]

# -1 means "infer this dimension"
A_reshaped = x.reshape((3, -1))  # 3 rows, infer columns
print(A_reshaped.shape)  # (3, 4)

Transpose

import jax.numpy as jnp

A = jnp.array([[1, 2, 3],
               [4, 5, 6]])
print(A.shape)      # (2, 3)

# Transpose
A_T = A.T
print(A_T.shape)    # (3, 2)
print(A_T)          # [[1 4]
                    #  [2 5]
                    #  [3 6]]

# For N-D arrays, specify axes
B = jnp.arange(24).reshape((2, 3, 4))
print(B.shape)      # (2, 3, 4)

# Swap axes 0 and 2
C = jnp.transpose(B, axes=(2, 1, 0))
print(C.shape)      # (4, 3, 2)

Squeeze and Expand

import jax.numpy as jnp

# Squeeze (remove dimensions of size 1)
x = jnp.array([[[1], [2]], [[3], [4]]])
print(x.shape)      # (2, 2, 1)

x_squeezed = jnp.squeeze(x)
print(x_squeezed.shape)  # (2, 2)

x_squeezed = jnp.squeeze(x, axis=2)  # Remove axis 2
print(x_squeezed.shape)  # (2, 2)

# Expand (add dimension)
x = jnp.array([1, 2, 3])
print(x.shape)      # (3,)

x_expanded = jnp.expand_dims(x, axis=0)
print(x_expanded.shape)  # (1, 3)

x_expanded = jnp.expand_dims(x, axis=1)
print(x_expanded.shape)  # (3, 1)

---

Concatenation and Stacking

Concatenate

import jax.numpy as jnp

# 1D arrays
x = jnp.array([1, 2, 3])
y = jnp.array([4, 5, 6])

z = jnp.concatenate([x, y])  # [1 2 3 4 5 6]

# 2D arrays
A = jnp.array([[1, 2], [3, 4]])
B = jnp.array([[5, 6], [7, 8]])

# Concatenate along axis 0 (rows)
C = jnp.concatenate([A, B], axis=0)
print(C.shape)  # (4, 2)

# Concatenate along axis 1 (columns)
C = jnp.concatenate([A, B], axis=1)
print(C.shape)  # (2, 4)

Stack

import jax.numpy as jnp

# Stack creates new axis
x = jnp.array([1, 2, 3])
y = jnp.array([4, 5, 6])

# Stack along new axis 0
z = jnp.stack([x, y], axis=0)
print(z.shape)  # (2, 3)
print(z)        # [[1 2 3]
                #  [4 5 6]]

# Stack along new axis 1
z = jnp.stack([x, y], axis=1)
print(z.shape)  # (3, 2)
print(z)        # [[1 4]
                #  [2 5]
                #  [3 6]]

---

Indexing and Slicing

Basic Indexing

import jax.numpy as jnp

x = jnp.array([0, 1, 2, 3, 4])

# Single index
print(x[0])     # 0
print(x[2])     # 2
print(x[-1])    # 4 (last element)

# Slicing
print(x[1:3])   # [1 2]
print(x[::2])   # [0 2 4] (every 2nd element)
print(x[::-1])  # [4 3 2 1 0] (reversed)

2D Indexing

import jax.numpy as jnp

A = jnp.array([[1, 2, 3],
               [4, 5, 6],
               [7, 8, 9]])

# Single element
print(A[0, 1])      # 2

# Row
print(A[1])         # [4 5 6]

# Column
print(A[:, 1])      # [2 5 8]

# Submatrix
print(A[1:3, :2])   # [[4 5]
                    #  [7 8]]

Fancy Indexing (Careful in JAX!)

import jax.numpy as jnp

x = jnp.array([10, 20, 30, 40, 50])

# NumPy fancy indexing
indices = jnp.array([0, 2, 4])
# print(x[indices])  # ❌ Might not work as expected in JAX!

# ✅ Use gather instead
result = jax.numpy.take(x, indices)
print(result)  # [10 30 50]

Updating Values (Immutable Update)

import jax.numpy as jnp

x = jnp.array([1, 2, 3, 4, 5])

# ❌ Can't do this in JAX:
# x[0] = 999  # TypeError!

# ✅ Use .at[] for immutable updates
x = x.at[0].set(999)
print(x)  # [999 2 3 4 5]

# Update multiple elements
x = jnp.array([1, 2, 3, 4, 5])
x = x.at[0, 2, 4](/0,-2,-4/).set(999)
print(x)  # [999 2 999 4 999]

# Add to elements
x = jnp.array([1, 2, 3, 4, 5])
x = x.at[0].add(10)
print(x)  # [11 2 3 4 5]

# Multiple operations
x = jnp.array([1, 2, 3, 4, 5])
x = x.at[0].add(10).at[2].multiply(100)
print(x)  # [11 2 300 4 5]

---

Linear Algebra

Matrix Operations

import jax.numpy as jnp

A = jnp.array([[1., 2.], [3., 4.]])
B = jnp.array([[5., 6.], [7., 8.]])
v = jnp.array([1., 2.])

# Matrix multiplication
C = jnp.dot(A, B)
print(C)  # [[19 22]
          #  [43 50]]

# Or use @
C = A @ B

# Matrix-vector multiplication
u = jnp.dot(A, v)
print(u)  # [5. 11.]

# Outer product
outer = jnp.outer(v, v)
print(outer)  # [[1 2]
              #  [2 4]]

# Inner product
inner = jnp.dot(v, v)
print(inner)  # 5.

Decompositions

import jax.numpy as jnp

A = jnp.array([[1., 2., 3.],
               [4., 5., 6.]])

# Singular Value Decomposition
U, S, Vt = jnp.linalg.svd(A, full_matrices=False)

# QR Decomposition
Q, R = jnp.linalg.qr(A)

# Eigenvalues and eigenvectors (square matrix)
A_square = jnp.array([[1., 2.], [2., 3.]])
eigenvalues, eigenvectors = jnp.linalg.eigh(A_square)

Matrix Inverse and Determinant

import jax.numpy as jnp

A = jnp.array([[1., 2.], [3., 4.]])

# Inverse
A_inv = jnp.linalg.inv(A)
print(A @ A_inv)  # ~identity matrix

# Determinant
det = jnp.linalg.det(A)
print(det)  # -2.

# Rank
rank = jnp.linalg.matrix_rank(A)
print(rank)  # 2

Solving Linear Systems

import jax.numpy as jnp

# Solve Ax = b
A = jnp.array([[1., 2.], [3., 4.]])
b = jnp.array([5., 6.])

x = jnp.linalg.solve(A, b)
print(x)  # Solution
print(A @ x)  # ~[5. 6.]

---

Broadcasting

Broadcasting automatically aligns arrays with different shapes.

Broadcasting Rules

import jax.numpy as jnp

# Scalar + Array
x = jnp.array([1, 2, 3])
y = x + 5  # Scalar broadcast to all elements
print(y)   # [6 7 8]

# Vector + Matrix
v = jnp.array([1, 2, 3])
A = jnp.array([[1, 2, 3],
               [4, 5, 6]])
B = A + v  # v broadcast to each row
print(B)   # [[2 4 6]
           #  [5 7 9]]

# Shape (3,) broadcasts to (2, 3)

Broadcasting Rules (Formal)

When operating on arrays: 1. Shapes aligned from right 2. Missing dimensions treated as size 1 3. Dimensions of size 1 broadcast to match other size

# (3,) broadcasts with (2, 3) → (2, 3)
#       Missing left dim treated as 1
#       (1, 3) broadcasts to (2, 3)

# (4, 1) broadcasts with (1, 3) → (4, 3)
#        1's broadcast to larger dimension

Practical Examples

import jax.numpy as jnp

# Subtract mean from each column
A = jnp.array([[1., 2., 3.],
               [4., 5., 6.]])
mean = jnp.mean(A, axis=0)  # [2.5 3.5 4.5] (shape 3,)
A_centered = A - mean       # Broadcast (3,) to (2, 3)

# Divide each column by std
std = jnp.std(A, axis=0)    # shape (3,)
A_normalized = A / std      # Broadcast (3,) to (2, 3)

# Outer subtraction
v1 = jnp.array([1, 2, 3])   # shape (3,)
v2 = jnp.array([10, 20])    # shape (2,)
# Reshape for broadcasting
diff = v1[None, :] - v2[:, None]  # (1, 3) - (2, 1) → (2, 3)

---

Performance Tips

Type Consistency

import jax.numpy as jnp

# ❌ Mixing dtypes forces conversion
x = jnp.array([1., 2., 3.], dtype=jnp.float32)
y = jnp.array([1., 2., 3.], dtype=jnp.float64)
z = x + y  # Converts to float64 (slower)

# ✅ Keep types consistent
z = x + x  # Both float32

Avoid Unnecessary Copies

import jax.numpy as jnp

x = jnp.array([1, 2, 3])

# ❌ Multiple copies
y = x.reshape(...).T.copy()

# ✅ Single operation
y = jnp.transpose(x.reshape(...))

---

Common Gotchas

Default dtype is float32

import jax.numpy as jnp

x = jnp.array([1., 2., 3.])
print(x.dtype)  # float32, not float64!

Indexing operations

import jax.numpy as jnp

x = jnp.array([1, 2, 3, 4, 5])

# In JAX, fancy indexing less flexible
# Use take, gather, scatter operations

Immutable updates

import jax.numpy as jnp

x = jnp.array([1, 2, 3])

# ❌ Can't do this
# x[0] = 999

# ✅ Must do this
x = x.at[0].set(999)

---

Summary Table

Operation	Code
Create array	jnp.array([1, 2, 3])
Reshape	x.reshape((3, 4))
Transpose	x.T or jnp.transpose(x)
Flatten	x.flatten()
Sum	jnp.sum(x, axis=0)
Mean	jnp.mean(x)
Min/Max	jnp.min(x), jnp.max(x)
Dot product	jnp.dot(x, y) or x @ y
Element-wise ops	x + y, x * y, jnp.sin(x)
Concatenate	jnp.concatenate([x, y])
Stack	jnp.stack([x, y])
Index	x[0], x[1:3], x[:, 1]
Update	x.at[0].set(5)

---

Next Steps

• 03 Pure Functions & Immutability - Master functional programming
• 01 Automatic Differentiation Jax.Grad - Learn gradients

Checkpoint: Explain the difference between .reshape() and .T (transpose)

---

Last Updated: 2026-08-09