[Python] Why is array manipulation in JAX much slower?

I'm working on converting a transformation-heavy numerical pipeline from NumPy to JAX to take advantage of JIT acceleration. However, I’ve found that some basic operations like broadcast_to and moveaxis are significantly slower in JAX—even without JIT—compared to NumPy, and even for large batch sizes like 3,000,000 where I would expect JAX to be much quicker.

### Benchmark: moveaxis + broadcast_to ###
NumPy: moveaxis + broadcast_to → 0.000116 s
JAX: moveaxis + broadcast_to → 0.204249 s
JAX JIT: moveaxis + broadcast_to → 0.054713 s

### Benchmark: broadcast_to only ###
NumPy: broadcast_to → 0.000059 s
JAX: broadcast_to → 0.062167 s
JAX JIT: broadcast_to → 0.057625 s


Am I doing something wrong? Are there better ways of performing these kind of manipulations?

Here's a minimal benchmark ChatGPT generated, comparing broadcast_to and moveaxis in NumPy, JAX, and JAX with JIT:

import timeit

import jax
import jax.numpy as jnp
import numpy as np
from jax import jit

# Base transformation matrix
M_np = np.array([[1, 0, 0, 0.5],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]])

M_jax = jnp.array(M_np)

# Batch size
n = 1_000_000

print("### Benchmark: moveaxis + broadcast_to ###")

# NumPy
t_numpy = timeit.timeit(
lambda: np.moveaxis(np.broadcast_to(M_np[:, :, None], (4, 4, n)), 2, 0),
number=10
)
print(f"NumPy: moveaxis + broadcast_to → {t_numpy:.6f} s")

# JAX
t_jax = timeit.timeit(
lambda: jnp.moveaxis(jnp.broadcast_to(M_jax[:, :, None], (4, 4, n)), 2, 0).block_until_ready(),
number=10
)
print(f"JAX: moveaxis + broadcast_to → {t_jax:.6f} s")

# JAX JIT
@jit
def broadcast_and_move_jax(M):
return jnp.moveaxis(jnp.broadcast_to(M[:, :, None], (4, 4, n)), 2, 0)

# Warm-up
broadcast_and_move_jax(M_jax).block_until_ready()

t_jit = timeit.timeit(
lambda: broadcast_and_move_jax(M_jax).block_until_ready(),
number=10
)
print(f"JAX JIT: moveaxis + broadcast_to → {t_jit:.6f} s")

print("\n### Benchmark: broadcast_to only ###")

# NumPy
t_numpy_b = timeit.timeit(
lambda: np.broadcast_to(M_np[:, :, None], (4, 4, n)),
number=10
)
print(f"NumPy: broadcast_to → {t_numpy_b:.6f} s")

# JAX
t_jax_b = timeit.timeit(
lambda: jnp.broadcast_to(M_jax[:, :, None], (4, 4, n)).block_until_ready(),
number=10
)
print(f"JAX: broadcast_to → {t_jax_b:.6f} s")

# JAX JIT
@jit
def broadcast_only_jax(M):
return jnp.broadcast_to(M[:, :, None], (4, 4, n))

broadcast_only_jax(M_jax).block_until_ready()

t_jit_b = timeit.timeit(
lambda: broadcast_only_jax(M_jax).block_until_ready(),
number=10
)
print(f"JAX JIT: broadcast_to → {t_jit_b:.6f} s")

Continue reading...