jax-ml/jax · error · ValueError

Index in subset_by_index[1] exceeds matrix size.

Error message

Index in subset_by_index[1] exceeds matrix size.

What it means

The end index of subset_by_index cannot exceed the matrix rank, i.e. min(m, n) of the input's last two dims. Only up to min(m, n) singular values exist, so the range must fit within that.

Source

Thrown at jax/_src/tpu/linalg/svd.py:190

      'specified to use `svd` within JAX transformations.',
  )

  if subset_by_index is not None:
    if len(subset_by_index) != 2:
      raise ValueError('subset_by_index must be a tuple of size 2.')
    # Make sure subset_by_index is a concrete tuple.
    subset_by_index = (
        operator.index(subset_by_index[0]),
        operator.index(subset_by_index[1]),
    )
    if subset_by_index[0] >= subset_by_index[1]:
      raise ValueError('Got empty index range in subset_by_index.')
    if subset_by_index[0] < 0:
      raise ValueError('Indices in subset_by_index must be non-negative.')
    m, n = a.shape
    rank = n if n < m else m
    if subset_by_index[1] > rank:
      raise ValueError('Index in subset_by_index[1] exceeds matrix size.')
    if full_matrices and subset_by_index != (0, rank):
      raise ValueError(
          'full_matrices and subset_by_index cannot be both be set.'
      )
    # By convention, eigenvalues are numbered in non-decreasing order, while
    # singular values are numbered non-increasing order, so change
    # subset_by_index accordingly.
    subset_by_index = (rank - subset_by_index[1], rank - subset_by_index[0])

  m, n = a.shape
  is_flip = False
  if m < n:
    a = a.T.conj()
    m, n = a.shape
    is_flip = True

  u_out_null: Array | None
  q: Array | None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp end: end = min(end, min(a.shape[-2], a.shape[-1])).
2. For full SVD, omit subset_by_index instead of passing (0, rank).
3. Add a shape assertion in data pipelines feeding matrices of varying sizes.

Example fix

// before
svd(a, subset_by_index=(0, 64))  # fails when min(m, n) < 64
// after
rank = min(a.shape[-2], a.shape[-1])
svd(a, subset_by_index=(0, min(64, rank)))

rank = min(a.shape[-2], a.shape[-1])
lo, hi = max(0, sb[0]), min(sb[1], rank)
assert lo < hi

def fits_rank(sb, a) -> bool:
    return sb[1] <= min(a.shape[-2], a.shape[-1])

try:
    svd(a, subset_by_index=sb)
except ValueError as e:
    if 'exceeds matrix size' in str(e): sb = (sb[0], min(a.shape[-1], a.shape[-2]))
    else: raise

When it happens

Trigger: svd(a, subset_by_index=(0, min(m, n) + 1)), or using the larger dimension instead of min(m, n) when computing the end bound, or running the same top-k code against smaller matrices than it was written for.

Common situations: Hard-coded k larger than the matrix's smaller dimension; batched code where some inputs have different shapes; assuming end can be max(m, n).

Related errors

• subset_by_index must be a tuple of size 2.
• Got empty index range in subset_by_index.
• Indices in subset_by_index must be non-negative.
• full_matrices and subset_by_index cannot be both be set.
• The SVD algorithm parameter is not implemented on TPU.

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