jax-ml/jax · error · ValueError

got {nse=}, expected to be between 0 and {sparse_size}

Error message

got {nse=}, expected to be between 0 and {sparse_size}

What it means

random_bcoo checks that nse (number of specified elements) lies in [0, sparse_size), where sparse_size is the product of the sparse dimensions. A fractional nse in (0,1) is interpreted as a density fraction; values >= sparse_size or negative are rejected.

Source

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

    generator : function for generating random values accepting a key, shape,
      and dtype. It defaults to :func:`jax.random.uniform`, and may be any
      function with a similar signature.
    **kwds : additional keyword arguments to pass to ``generator``.

  Returns:
    arr : a sparse.BCOO array with the specified properties.
  """
  shape = tuple(map(operator.index, shape))
  n_batch = operator.index(n_batch)
  n_dense = operator.index(n_dense)
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass nse < math.prod(sparse_dims), or use a fractional density like 0.1 for sparse random matrices
  2. For a dense result use random.normal/keyed dense generation instead
  3. Recompute nse = int(density * sparse_size) with density < 1

Example fix

// before
M = random_bcoo(key, shape=(4, 4), nse=16)
// after
M = random_bcoo(key, shape=(4, 4), nse=0.5)  # ~8 elements
Defensive patterns

Strategy: validation

Validate before calling

sparse_size = math.prod(shape[n_batch:len(shape)-n_dense])
assert 0 <= nse < sparse_size, f'nse={nse} out of range for sparse size {sparse_size}'

Prevention

When it happens

Trigger: Passing an absolute nse larger than or equal to the number of sparse elements (e.g. nse=16 for a 4x4 matrix), or a negative nse, or a density fraction >= 1.

Common situations: Treating nse as a density (passing 0.5 works but 1.0 fails); computing nse from a bigger matrix shape than the one passed; off-by-one when nse == sparse_size for a 'fully dense' sparse matrix.

Related errors


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