jax-ml/jax · error · ValueError

{indices_dtype=} does not have enough range to generate spar

Error message

{indices_dtype=} does not have enough range to generate sparse indices of size {sparse_size}.

What it means

random_bcoo generates flat sparse indices with random.choice over the sparse extent; if the chosen indices_dtype (default int32) cannot represent sparse_size-1, index generation would overflow, so it raises preemptively.

Source

Thrown at jax/experimental/sparse/random.py:87

  if n_batch < 0 or n_dense < 0 or n_batch + n_dense > len(shape):
    raise ValueError(f"Invalid {n_batch=}, {n_dense=} for {shape=}")
  n_sparse = len(shape) - n_batch - n_dense
  batch_shape, sparse_shape, dense_shape = map(tuple, split_list(shape, [n_batch, n_sparse]))
  batch_size = math.prod(batch_shape)
  sparse_size = math.prod(sparse_shape)
  if not 0 <= nse < sparse_size:
    raise ValueError(f"got {nse=}, expected to be between 0 and {sparse_size}")
  if 0 < nse < 1:
    nse = int(math.ceil(nse * sparse_size))
  assert not isinstance(nse, float)
  nse = operator.index(nse)

  data_shape = batch_shape + (nse,) + dense_shape
  indices_shape = batch_shape + (nse, n_sparse)
  if indices_dtype is None:
    indices_dtype = dtypes.default_int_dtype()
  if sparse_size > jnp.iinfo(indices_dtype).max:
    raise ValueError(f"{indices_dtype=} does not have enough range to generate "
                     f"sparse indices of size {sparse_size}.")
  @vmap
  def _indices(key):
    if not sparse_shape:
      return jnp.zeros((nse, n_sparse), dtype=indices_dtype)
    flat_ind = random.choice(key, sparse_size, shape=(nse,),
                             replace=not unique_indices).astype(indices_dtype)
    return jnp.column_stack(jnp.unravel_index(flat_ind, sparse_shape))

  keys = random.split(key, batch_size + 1)
  data_key, index_keys = keys[0], keys[1:]
  data = generator(data_key, shape=data_shape, dtype=dtype, **kwds)
  indices = _indices(index_keys).reshape(indices_shape)
  mat = sparse.BCOO((data, indices), shape=shape)
  return mat.sort_indices() if sorted_indices else mat

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass indices_dtype=jnp.int64 (and enable 64-bit ints via jax_enable_x64 if needed)
Reduce the sparse dimensions' product below the dtype's max
Leave indices_dtype=None to use the default int dtype, which must still be large enough

Example fix

// before
M = random_bcoo(key, shape=(10**6, 10**6), nse=100)
// after
jax.config.update('jax_enable_x64', True)
M = random_bcoo(key, shape=(10**6, 10**6), nse=100, indices_dtype=jnp.int64)

Defensive patterns

Strategy: validation

Validate before calling

import math, jax.numpy as jnp
sparse_size = math.prod(shape)
if sparse_size > jnp.iinfo(indices_dtype or jnp.int32).max:
    jax.config.update('jax_enable_x64', True)
    indices_dtype = jnp.int64

Prevention

Enable 64-bit mode and pass indices_dtype=jnp.int64 for very large sparse layouts
Estimate product of sparse dims against the dtype max before generation

When it happens

Trigger: Generating a BCOO whose sparse dimensions multiply to more than jnp.iinfo(indices_dtype).max — e.g. a 100000 x 100000 sparse layout with default int32 indices.

Common situations: Very large sparse matrices with default int32; explicitly requesting int8/int16 indices; enabling x64 only for data but not indices.

Related errors

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