jax-ml/jax · error · ValueError

subset_by_index must be a tuple of size 2.

Error message

subset_by_index must be a tuple of size 2.

What it means

jax._src.tpu.linalg.svd raises this when the subset_by_index argument passed to jax.scipy.linalg.svd (TPU path) is not a 2-element tuple. subset_by_index selects a contiguous range of singular values/vectors to compute, and it must be specified as exactly (start, end) with concrete Python ints.

Source

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

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

  hermitian = core.concrete_or_error(
      bool,
      hermitian,
      'The `hermitian` argument must be statically '
      'specified to use `svd` within JAX transformations.',
  )

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a 2-tuple of concrete ints: subset_by_index=(start, end), end exclusive, 0 <= start < end <= min(m, n).
  2. If you don't need a partial SVD, omit subset_by_index entirely.
  3. If computing bounds dynamically, assert len == 2 and both are Python ints (use operator.index) before calling.

Example fix

// before
u, s, vt = svd(a, subset_by_index=(k,))
// after
u, s, vt = svd(a, subset_by_index=(0, k))
Defensive patterns

Strategy: validation

Validate before calling

from operator import index
def valid_subset(sb, a):
    rank = min(a.shape[-2], a.shape[-1])
    assert sb is None or (len(sb) == 2 and all(isinstance(index(i), int) for i in sb)), 'need (start, end)'
    return sb

Type guard

def is_subset_by_index(x) -> bool:
    return x is None or (isinstance(x, tuple) and len(x) == 2 and all(isinstance(v, int) for v in x))

Try / catch

try:
    u, s, vt = svd(a, subset_by_index=sb)
except ValueError as e:
    if 'tuple of size 2' in str(e): raise ValueError(f'bad subset_by_index: {sb!r}')
    raise

Prevention

When it happens

Trigger: Calling svd(a, subset_by_index=(0,)) or subset_by_index=[0, 3, 5] — any value whose len() != 2, e.g. passing a single int or a 3-tuple.

Common situations: Confusing subset_by_index with numpy's subset_by_index style APIs, or computing the tuple dynamically (e.g. from array shape arithmetic) and accidentally producing a 1- or 3-element sequence.

Related errors


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