jax-ml/jax · error · ValueError

Indices in subset_by_index must be non-negative.

Error message

Indices in subset_by_index must be non-negative.

What it means

Both elements of subset_by_index must be non-negative integers. Negative indices (Python-style negative indexing) are not supported for selecting a singular value range in the TPU SVD.

Source

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

  max_iterations = core.concrete_or_error(
      int,
      max_iterations,
      'The `max_iterations` argument must be statically '
      '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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Normalize negative indices to positive: start = start % rank before calling.
2. Use only 0-based non-negative bounds: (0, k) for top-k.
3. Validate bounds against matrix shape min(m, n) before the call.

Example fix

// before
svd(a, subset_by_index=(-k, rank))
// after
svd(a, subset_by_index=(rank - k, rank))

if sb[0] < 0:
    rank = min(a.shape[-2], a.shape[-1])
    sb = (sb[0] % rank, sb[1])  # or: (rank + sb[0], rank)

def has_nonneg_bounds(sb) -> bool:
    return all(i >= 0 for i in sb)

try:
    svd(a, subset_by_index=sb)
except ValueError as e:
    if 'non-negative' in str(e): sb = tuple(i % rank for i in sb)
    else: raise

When it happens

Trigger: svd(a, subset_by_index=(-2, 3)) or any start < 0, e.g. porting numpy-style a[-3:] indexing logic into subset_by_index.

Common situations: Developers assuming numpy negative-index semantics; range computed as min(k, n-k) minus a constant that can dip below zero for small matrices.

Related errors

• subset_by_index must be a tuple of size 2.
• Got empty index range in subset_by_index.
• Index in subset_by_index[1] exceeds matrix size.
• 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/aa04530d97e23fa5. Report an issue: GitHub.