jax-ml/jax · error · ValueError

bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndi

Error message

bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndim}

What it means

bcoo_dynamic_slice validates that len(start_indices) == len(slice_sizes) == mat.ndim via a chained comparison; any mismatch raises ValueError. start_indices must be one scalar index per dimension and slice_sizes one size per dimension.

Source

Thrown at jax/experimental/sparse/bcoo.py:2062

      integers with length equal to `ndim(operand)`. Inside a JIT compiled
      function, only static values are supported (all JAX arrays inside JIT
      must have statically known size).

  Returns:
    out: BCOO array containing the slice.
  """
  slice_sizes = tuple(operator.index(i) for i in slice_sizes)
  # Use abstract eval to validate inputs.
  jax.jit(lax.dynamic_slice, static_argnames=("slice_sizes",)).eval_shape(
          jax.ShapeDtypeStruct(mat.shape, mat.dtype), start_indices,
          slice_sizes=slice_sizes)
  if not isinstance(mat, BCOO):
    raise TypeError(f"bcoo_slice: input should be BCOO array, got type(mat)={type(mat)}")
  start_indices = tuple(jnp.asarray(i) for i in start_indices)
  assert all(jnp.issubdtype(i.dtype, np.integer) for i in start_indices)
  assert all(i.shape == () for i in start_indices)
  if len(start_indices) != len(slice_sizes) != mat.ndim:
    raise ValueError(f"bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndim}")
  if not all(0 <= slice_size <= axis_size for slice_size, axis_size in zip(slice_sizes, mat.shape)):
    raise TypeError("slice_sizes must be less than or equal to operand shape, "
                    f"got slice_sizes {slice_sizes} for operand shape {mat.shape}")

  start_batch, start_sparse, start_dense = split_list(start_indices, [mat.n_batch, mat.n_sparse])
  size_batch, size_sparse, size_dense = split_list(slice_sizes, [mat.n_batch, mat.n_sparse])

  data_start = []
  data_sizes = []
  indices_start = []
  indices_sizes = []
  zero = _const(start_indices[0] if start_indices else np.int32, 0)
  for i, (start, size) in enumerate(zip(start_batch, size_batch)):
    data_is_broadcast = mat.data.shape[i] != mat.shape[i]
    indices_is_broadcast = mat.indices.shape[i] != mat.shape[i]
    data_start.append(zero if data_is_broadcast else start)
    data_sizes.append(1 if data_is_broadcast else size)
    indices_start.append(zero if indices_is_broadcast else start)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Provide exactly mat.ndim scalar start indices and mat.ndim slice sizes
Construct the tuples from mat.ndim at runtime rather than hardcoding
Note the lax-level eval_shape validation also runs first — fix lengths so both checks pass

Example fix

# before (mat is 2-D)
bcoo_dynamic_slice(mat, start_indices=(i,), slice_sizes=(4,))
# after
bcoo_dynamic_slice(mat, start_indices=(i, 0), slice_sizes=(4, 4))

Defensive patterns

Strategy: validation

Validate before calling

start_indices = tuple(start_indices)
slice_sizes = tuple(slice_sizes)
assert len(start_indices) == len(slice_sizes) == mat.ndim

Prevention

Build start/size tuples from mat.ndim, not fixed literals
Remember batch dims count toward ndim in the index lists

When it happens

Trigger: Calling bcoo_dynamic_slice with fewer/more start indices or slice sizes than mat.ndim, e.g. a single start index and size for a 2-D BCOO array.

Common situations: Porting lax.dynamic_slice code between arrays of different rank; assuming batch dimensions are excluded from the index lists; building slice_sizes dynamically with the wrong length.

Related errors

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