jax-ml/jax · error · ValueError

jax.numpy.block does not allow tuples, got {}

Error message

jax.numpy.block does not allow tuples, got {}

What it means

jax.numpy.block accepts only lists (not tuples) as the nested layout specification, unlike NumPy which accepts either. The internal _block helper explicitly rejects tuple inputs with this ValueError. This is a deliberate JAX API restriction for consistency.

Source

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

      raise ValueError("invalid entry in choice array")
  elif mode == 'wrap':
    arr = asarray(a) % N
  elif mode == 'clip':
    arr = clip(a, 0, N - 1)
  else:
    raise ValueError(f"mode={mode!r} not understood. Must be 'raise', 'wrap', or 'clip'")

  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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Replace all parentheses with square brackets: jnp.block([[a, b], [c, d]])
  2. If the structure comes from a tuple-producing pipeline, convert it: jnp.block(list(map(list, nested_tuple)))

Example fix

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

Strategy: validation

Validate before calling

xs = list(map(list, xs)) if isinstance(xs, tuple) else xs
out = jnp.block(xs)

Type guard

def is_block_list(xs) -> bool:
    return isinstance(xs, list)

Prevention

When it happens

Trigger: Calling jnp.block(((a, b), (c, d))) with parentheses instead of brackets — i.e. tuples of tuples — anywhere in the nesting, including a single level like jnp.block((a, b)).

Common situations: Code ported from np.block that used tuples; or automatic conversion of lists to tuples by libraries like dataclasses, namedtuples, or pytrees that then feed the result into jnp.block.

Related errors


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