`jnp.concatenate` on minor axes significantly slower in newer versions

Description

Summary

Running jnp.concatenate(..., axis=1) takes a very long time to run. This wasn't the case in jax<0.11.0, but is the case in jax>=0.11.0. I understand that there may be some optimizations for concatenating along axis=0, and indeed, this is a bit faster in jax<0.11.0, but things become much slower in jax>0.11.0 (potentially taking hours for a few thousand elements).

I'm able to reproduce this in a fresh Colab notebook.

Benchmarking

Here's some concise benchmarking code (in the Colab notebook):

import platform, time
import jax, jax.numpy as jnp

print(f"jax {jax.__version__}   backend {jax.default_backend()}   {platform.platform()}")
WIDTH = 400

def make(n, shape):
    xs = [jnp.zeros(shape) for _ in range(n)]
    jax.block_until_ready(xs)
    return xs

def compile_s(xs, axis):
    lowered = jax.jit(lambda ys: jnp.concatenate(ys, axis=axis)).lower(xs)
    t0 = time.perf_counter(); lowered.compile()
    return time.perf_counter() - t0

def eager_s(xs, axis):
    t0 = time.perf_counter()
    jnp.concatenate(xs, axis=axis).block_until_ready()
    return time.perf_counter() - t0

compile_s(make(2, (WIDTH, 1)), 1)  # warm up

print("\n             axis=0 control            axis=1 (regressed)")
print("      N   (1,400)xN -> (N,400)      (400,1)xN -> (400,N)")
print("            compile      eager        compile      eager")
for n in (50, 100, 200, 400, 800, 1600):
    major, minor = make(n, (1, WIDTH)), make(n, (WIDTH, 1))
    c0, e0 = compile_s(major, 0), eager_s(major, 0)
    c1, e1 = compile_s(minor, 1), eager_s(minor, 1)
    print(f"{n:>7}   {c0:8.3f}s  {e0:8.3f}s     {c1:9.3f}s  {e1:9.3f}s")
    if c1 > 60:
        print("          (stopping: axis=1 compile exceeds 60s timeout)")
        break

JAX 0.10.x Results

The results for jax<0.11.0 give that the minor axis concatenation is somewhat slower than major access concatenation, though bearably so:

jax 0.10.2   backend cpu   Linux-6.6.122+-x86_64-with-glibc2.39

             axis=0 control            axis=1 (regressed)
      N   (1,400)xN -> (N,400)      (400,1)xN -> (400,N)
            compile      eager        compile      eager
     50      0.096s     0.203s         0.098s      0.264s
    100      0.304s     0.199s         0.364s      0.251s
    200      0.281s     0.290s         0.963s      0.584s
    400      0.278s     0.298s         1.180s      1.201s
    800      0.804s     0.827s         5.804s      4.519s
   1600      2.960s     3.985s        21.745s     22.511s

JAX 0.11.x Results

Compile along major axis is comparable, but compile along minor axis with $N=800$ goes from 5.8s -> >300s.

jax 0.11.1   backend cpu   Linux-6.6.122+-x86_64-with-glibc2.39

             axis=0 control            axis=1 (regressed)
      N   (1,400)xN -> (N,400)      (400,1)xN -> (400,N)
            compile      eager        compile      eager
     50      0.165s     0.142s         0.207s      0.321s
    100      0.343s     0.381s         0.375s      0.423s
    200      0.400s     0.380s         6.563s      5.374s
    400      0.360s     0.327s        35.836s     37.034s
    800      0.831s     1.631s       313.703s    313.004s
          (stopping: axis=1 compile exceeds 60s timeout)

This scaling looks ~cubic, and can take hours with a few thousand elements.

AI Use

For the sake of AI use transparency: I ran into this bug "in the wild" (cf. https://github.com/BasisResearch/dynestyx/issues/353), and used Claude to debug and give minimal benchmarking code. I wrote this issue myself.

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.11.1
jaxlib: 0.11.1
numpy:  2.1.3
python: 3.13.15 (main, Aug  6 2026, 11:06:22) [GCC 13.3.0]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='dff9cf821d76', release='6.6.122+', version='#1 SMP Thu Apr 30 18:17:14 UTC 2026', machine='x86_64')