jax-ml/jax · error · ValueError

Invalid {n_batch=}, {n_dense=} for {shape=}

Error message

Invalid {n_batch=}, {n_dense=} for {shape=}

What it means

random_bcoo validates the layout split of the output shape: n_batch (batch dims), n_dense (dense/trailing dims) must each be >= 0 and their sum must not exceed the total rank of shape.

Source

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

    n_dense : number of batch dimensions. must satisfy ``n_dense >= 0`` and
      ``n_batch + n_dense <= len(shape)``.
    unique_indices : boolean specifying whether indices should be unique
      (default: True).
    sorted_indices : boolean specifying whether indices should be row-sorted in
      lexicographical order (default: False).
    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}.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure 0 <= n_batch and 0 <= n_dense and n_batch + n_dense <= len(shape)
2. For an all-sparse layout pass n_batch=0, n_dense=0
3. Recompute the counts from the final shape right before the call

Example fix

// before
M = random_bcoo(key, shape=(4, 4), n_batch=1, n_dense=4)
// after
M = random_bcoo(key, shape=(4, 4), n_batch=0, n_dense=0)

assert n_batch >= 0 and n_dense >= 0 and n_batch + n_dense <= len(shape)

When it happens

Trigger: Calling jax.experimental.sparse.random_bcoo(key, shape, n_batch, n_dense) with negative counts or n_batch + n_dense > len(shape).

Common situations: Computing n_dense = len(shape) - n_batch and then also passing n_batch, double-counting; rank changes after refactoring shape; passing len(shape) as n_dense for a fully-dense layout.

Related errors

• bcoo_slice: indices must have size mat.ndim={mat.ndim}
• len(strides) = {len(strides)}; expected {mat.ndim}
• bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndi
• got {nse=}, expected to be between 0 and {sparse_size}
• {indices_dtype=} does not have enough range to generate spar

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