jax-ml/jax · error · ValueError

Invalid dimension range passed to collapse: {operand.shape}[

Error message

Invalid dimension range passed to collapse: {operand.shape}[{start_dimension}:{stop_dimension}]

What it means

jax.lax.collapse ravels dimensions [start, stop) of an array into one. Python slice semantics clamp the range, and if the resolved stop index is below the start (hi < lo), the range is empty/negative in an unsupported way and JAX raises this ValueError.

Source

Thrown at jax/_src/lax/lax.py:3987

  For example, if ``operand`` is an array with shape ``[2, 3, 4]``,
  ``collapse(operand, 0, 2).shape == [6, 4]``. The elements of the collapsed
  dimension are laid out major-to-minor, i.e., with the lowest-numbered
  dimension as the slowest varying dimension.

  Args:
    operand: an input array.
    start_dimension: the start of the dimensions to collapse (inclusive).
    stop_dimension: the end of the dimensions to collapse (exclusive). Pass None
      to collapse all the dimensions after start.

  Returns:
    An array where dimensions ``[start_dimension, stop_dimension)`` have been
    collapsed (raveled) into a single dimension.
  """
  lo, hi, _ = slice(start_dimension, stop_dimension).indices(len(operand.shape))
  if hi < lo:
    raise ValueError(f"Invalid dimension range passed to collapse: {operand.shape}"
                     f"[{start_dimension}:{stop_dimension}]")
  size = math.prod(operand.shape[lo:hi])
  new_shape = operand.shape[:lo] + (size,) + operand.shape[hi:]
  return reshape(operand, new_shape)


def batch_matmul(lhs: Array, rhs: Array,
                 precision: PrecisionLike = None) -> Array:
  """Batch matrix multiplication."""
  if _min(lhs.ndim, rhs.ndim) < 2:
    raise ValueError('Arguments to batch_matmul must be at least 2D, got {}, {}'
                     .format(lhs.ndim, rhs.ndim))
  if lhs.ndim != rhs.ndim:
    raise ValueError('Arguments to batch_matmul must have same ndim, got {}, {}'
                     .format(lhs.ndim, rhs.ndim))
  lhs_contract = (lhs.ndim - 1,)
  rhs_contract = (rhs.ndim - 2,)
  batch = tuple(range(lhs.ndim - 2))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use positive absolute indices: for 'to the end' pass operand.ndim (or None-style equivalent by computing it)
2. Validate start <= stop and both within [0, x.ndim] before calling
3. Check your slice arithmetic (stop - start should be >= 0 after resolution)

Example fix

// before
collapsed = lax.collapse(x, 1, -1)
// after
collapsed = lax.collapse(x, 1, x.ndim)

start, stop = int(start), int(stop)
assert 0 <= start <= stop <= x.ndim, f'bad collapse range {start}:{stop}'
out = lax.collapse(x, start, stop)

When it happens

Trigger: Calling lax.collapse(x, start, stop) where stop normalizes before start, e.g. collapse(x, 2, 1), collapse(x, 0, -1) on small arrays, or stop computed as start - k by mistake.

Common situations: Passing negative stop_dimension expecting exclusive positive semantics; computing stop = start + k with a negative k; using -1 as 'until the end' (which is not how collapse works — omit it or pass x.ndim instead).

Related errors

• repeated axis in lax.expand_dims: {dimensions}
• Arguments to batch_matmul must be at least 2D, got {}, {}
• Arguments to batch_matmul must have same ndim, got {}, {}
• lax.platform_dependent: the '{pname}' branch must be a calla
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.

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