jax-ml/jax · error · ValueError

Arguments to jax.scipy.linalg.block_diag must have at most 2

Error message

Arguments to jax.scipy.linalg.block_diag must have at most 2 dimensions, got {} at argument {}.

What it means

jax.scipy.linalg.block_diag builds a block-diagonal matrix from its arguments and requires each to be at most 2-D (scalars, vectors, or matrices). Any argument with ndim > 2 raises ValueError, reporting the offending array and its position in the argument list.

Source

Thrown at jax/_src/scipy/linalg.py:1658

  Examples:
    >>> A = jnp.ones((1, 1))
    >>> B = jnp.ones((2, 2))
    >>> C = jnp.ones((3, 3))
    >>> jax.scipy.linalg.block_diag(A, B, C)
    Array([[1., 0., 0., 0., 0., 0.],
           [0., 1., 1., 0., 0., 0.],
           [0., 1., 1., 0., 0., 0.],
           [0., 0., 0., 1., 1., 1.],
           [0., 0., 0., 1., 1., 1.],
           [0., 0., 0., 1., 1., 1.]], dtype=float32)
  """
  if len(arrs) == 0:
    arrs = (jnp.zeros((1, 0)),)
  arrs = tuple(promote_dtypes(*arrs))
  bad_shapes = [i for i, a in enumerate(arrs) if np.ndim(a) > 2]
  if bad_shapes:
    raise ValueError("Arguments to jax.scipy.linalg.block_diag must have at "
                     "most 2 dimensions, got {} at argument {}."
                     .format(arrs[bad_shapes[0]], bad_shapes[0]))
  converted_arrs = [jnp.atleast_2d(a) for a in arrs]
  dtype = lax.dtype(converted_arrs[0])
  total_cols = sum(a.shape[1] for a in converted_arrs)

  padded_arrs = []
  current_col = 0
  for arr in converted_arrs:
    cols = arr.shape[1]
    padding_config = ((0, 0, 0), (current_col, total_cols - cols - current_col, 0))
    padded_arrs.append(lax.pad(arr, dtype.type(0), padding_config))
    current_col += cols
  return jnp.concatenate(padded_arrs, axis=0)


@jit(static_argnames=("eigvals_only", "select", "select_range"))
def eigh_tridiagonal(d: ArrayLike, e: ArrayLike, *, eigvals_only: bool = False,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Squeeze/reshape the offending argument to 2-D, e.g. arr = arr.reshape(-1, arr.shape[-1]) or arr[0] if batched
Check the argument index reported in the message to find which input is over-dimensional

Example fix

// before
blocks = jnp.stack([m1, m2])  # (2, 2, 2)
bd = jax.scipy.linalg.block_diag(blocks, m3)
// after
bd = jax.scipy.linalg.block_diag(m1, m2, m3)

Defensive patterns

Strategy: validation

Validate before calling

bad = [i for i, a in enumerate(arrs) if np.ndim(a) > 2]
assert not bad, f'args {bad} exceed 2 dims'

Type guard

def all_at_most_2d(arrs): return all(np.ndim(a) <= 2 for a in arrs)

Prevention

Unstack batches before block_diag; vmap block_diag if you need per-sample block diagonals
Check ndim of stacked intermediates (result of jnp.stack) before passing

When it happens

Trigger: Calling block_diag(a, b) where one argument is 3-D, e.g. a stacked batch of shape (4, 2, 2) from vmap or stacking.

Common situations: Feeding an accidentally stacked/extra-dimension array (e.g. result of jnp.stack or an un-squeezed matrix) into block_diag; SciPy has the same 2-D limit, so ported code with shape bugs surfaces here.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/907def87e0c8e505. Report an issue: GitHub.