Reshaping Arrays

Reshaping arrays is a constant operation in ML — converting between batch formats, flattening for dense layers, adding dimensions for broadcasting. JAX provides the same reshaping tools as NumPy.

import jax.numpy as jnp

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

# Reshape: change shape without changing data
a = jnp.reshape(x, (3, 4))   # 3 rows, 4 columns
b = x.reshape(3, 4)           # Method syntax works too
print(a)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

# -1 infers one dimension
c = x.reshape(2, -1)  # (2, 6) — JAX figures out 6
d = x.reshape(-1, 3)  # (4, 3)

Adding and Removing Dimensions

import jax.numpy as jnp

v = jnp.array([1.0, 2.0, 3.0])  # shape: (3,)

# expand_dims: add a dimension
row = jnp.expand_dims(v, axis=0)   # shape: (1, 3)
col = jnp.expand_dims(v, axis=1)   # shape: (3, 1)

# Equivalent using None/newaxis indexing
row2 = v[None, :]   # shape: (1, 3)
col2 = v[:, None]   # shape: (3, 1)

# squeeze: remove dimensions of size 1
x = jnp.zeros((1, 3, 1, 4))
squeezed = jnp.squeeze(x)          # shape: (3, 4)
partial = jnp.squeeze(x, axis=0)   # shape: (3, 1, 4)

Transposing and Permuting Axes

import jax.numpy as jnp

# 2D transpose
a = jnp.array([[1, 2, 3], [4, 5, 6]])
print(a.T.shape)  # (3, 2)

# Higher-dimensional: permute axes
# Common in ML: converting between channels-first and channels-last
img = jnp.zeros((32, 3, 224, 224))  # (batch, channels, height, width)

# NCHW -> NHWC
img_nhwc = jnp.transpose(img, (0, 2, 3, 1))
print(img_nhwc.shape)  # (32, 224, 224, 3)

Concatenation and Stacking

import jax.numpy as jnp

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

# Concatenate: join along existing axis
c = jnp.concatenate([a, b])
print(c)  # [1 2 3 4 5 6]

# Stack: join along a NEW axis
s = jnp.stack([a, b])
print(s)        # [[1 2 3], [4 5 6]]
print(s.shape)  # (2, 3)

# vstack and hstack
v = jnp.vstack([a, b])  # Same as stack for 1D → 2D
h = jnp.hstack([a, b])  # Same as concatenate for 1D