jax-ml/jax · error · ValueError

Mismatched list depths in jax.numpy.block

Error message

Mismatched list depths in jax.numpy.block

What it means

jnp.block requires the nested list structure to be a proper rectangular grid: every sub-list at the same level must have the same nesting depth. _block computes depths recursively and raises this ValueError when sibling depths differ, e.g. mixing a scalar with a list-of-lists.

Source

Thrown at jax/_src/numpy/lax_numpy.py:4999

  arr, *choices = broadcast_arrays(arr, *choices)
  return array(choices)[(arr,) + indices(arr.shape, sparse=True)]


def _atleast_nd(x: ArrayLike, n: int) -> Array:
  m = np.ndim(x)
  return lax.broadcast(x, (1,) * (n - m)) if m < n else asarray(x)

def _block(xs: ArrayLike | list[Any]) -> tuple[Array, int]:
  if isinstance(xs, tuple):
    raise ValueError("jax.numpy.block does not allow tuples, got {}"
                     .format(xs))
  elif isinstance(xs, list):
    if len(xs) == 0:
      raise ValueError("jax.numpy.block does not allow empty list arguments")
    xs_tup, depths = unzip2([_block(x) for x in xs])
    if any(d != depths[0] for d in depths[1:]):
      raise ValueError("Mismatched list depths in jax.numpy.block")
    rank = max(depths[0], max(np.ndim(x) for x in xs_tup))
    xs_tup = tuple(_atleast_nd(x, rank) for x in xs_tup)
    return concatenate(xs_tup, axis=-depths[0]), depths[0] + 1
  else:
    return asarray(xs), 1


@export
@api.jit
def block(arrays: ArrayLike | list[Any]) -> Array:
  """Create an array from a list of blocks.

  JAX implementation of :func:`numpy.block`.

  Args:
    arrays: an array, or nested list of arrays which will be concatenated
      together to form the final array.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap scalar/single-array rows in their own list: [[a, b], [c]] becomes [[a, b], [c, ]]... ensure consistent nesting
2. Validate depths before calling: all(len(r) == len(rows[0]) and isinstance(r, list) for r in rows)
3. Assemble via explicit jnp.concatenate calls per axis instead of block

Example fix

// before
out = jnp.block([[a, b], c])
// after
out = jnp.block([[a, b], [c]])

def uniform_depth(xs):
    if not isinstance(xs, list):
        return 0
    depths = {uniform_depth(x) for x in xs}
    return 1 + depths.pop() if len(depths) == 1 else -1
assert uniform_depth(blocks) >= 0

When it happens

Trigger: jnp.block([[a, b], c]) — second row is a raw array while the first is a list of two; or jnp.block([a, [b, c]]) where one element is a leaf and the other a list.

Common situations: Heterogeneous construction where some rows are single blocks (not wrapped in their own list) and others are multiple blocks; ported NumPy code that happened to work because shapes still concatenated correctly.

Related errors

• jax.numpy.block does not allow tuples, got {}
• jax.numpy.block does not allow empty list arguments
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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