jax-ml/jax · error · TypeError

slice_sizes must be less than or equal to operand shape, got

Error message

slice_sizes must be less than or equal to operand shape, got slice_sizes {slice_sizes} for operand shape {mat.shape}

What it means

bcoo_dynamic_slice requires each slice_sizes[i] to satisfy 0 <= slice_sizes[i] <= mat.shape[i]; a window larger than its axis raises TypeError (matching lax.dynamic_slice's contract). The lax abstract-eval check may also surface the same problem first.

Source

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

      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)
    indices_sizes.append(1 if indices_is_broadcast else size)
  data_start.append(zero)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Clamp each size: size = max(0, min(size, mat.shape[axis]))
  2. Compute sizes from actual shape at runtime instead of constants
  3. If you need clamped tail behavior, adjust start and size together before the call

Example fix

# before
bcoo_dynamic_slice(mat, (i,), (16,))  # axis size 10
# after
bcoo_dynamic_slice(mat, (i,), (min(16, mat.shape[0] - i),))
Defensive patterns

Strategy: validation

Validate before calling

slice_sizes = tuple(max(0, min(s, dim - start))
                     for s, dim, start in zip(slice_sizes, mat.shape, start_indices))

Try / catch

try:
    out = bcoo_dynamic_slice(mat, starts, sizes)
except TypeError as e:
    if 'slice_sizes' in str(e):
        sizes = tuple(min(s, d) for s, d in zip(sizes, mat.shape))
        out = bcoo_dynamic_slice(mat, starts, sizes)
    else:
        raise

Prevention

When it happens

Trigger: Requesting a window of size larger than the corresponding dimension, e.g. slice_sizes=(8,) on an axis of size 4, or negative sizes.

Common situations: Hardcoded window sizes applied to inputs of varying shapes; computing sizes as limit - start and getting negatives; assuming slicing past the end clamps like NumPy (it does not).

Related errors

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