tracel-ai/burn · error

Invalid dimension: the shape of the index tensor should be t

Error message

Invalid dimension: the shape of the index tensor should be the same as the value tensor: Index {:?} value {:?}

What it means

The NdArray `scatter` op panics when the index tensor's shape does not exactly equal the value tensor's shape. scatter (burn-ndarray/src/ops/base.rs) requires index/value shape equality by construction — every element of `values` has a corresponding index — so mismatched shapes are rejected with a panic naming both shapes.

Source

Thrown at crates/burn-ndarray/src/ops/base.rs:163

    ) -> SharedArray<E> {
        let ndims = tensor.shape().num_dims();
        if dim != ndims - 1 {
            tensor.swap_axes(ndims - 1, dim);
            indices.swap_axes(ndims - 1, dim);
            value.swap_axes(ndims - 1, dim);
        }

        let (shape_tensor, shape_indices, shape_value) =
            (tensor.shape().into_shape(), indices.shape(), value.shape());
        let (size_tensor, size_index, size_value) = (
            shape_tensor[ndims - 1],
            shape_indices[ndims - 1],
            shape_value[ndims - 1],
        );
        let batch_size = Self::gather_batch_size(&shape_tensor, shape_indices);

        if shape_value != shape_indices {
            panic!(
                "Invalid dimension: the shape of the index tensor should be the same as the value \
                 tensor: Index {:?} value {:?}",
                shape_indices, shape_value
            );
        }

        let indices = NdArrayOps::reshape(indices, Shape::new([batch_size, size_index]));
        let value = NdArrayOps::reshape(value, Shape::new([batch_size, size_value]));
        let mut tensor = NdArrayOps::reshape(tensor, Shape::new([batch_size, size_tensor]));

        for b in 0..batch_size {
            let indices = indices.slice(s!(b, ..));

            for (i, index) in indices.iter().enumerate() {
                let index = index.elem::<i64>() as usize;
                tensor[[b, index]].add_assign(value[[b, i]]);
            }
        }

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Reshape/expand the index tensor to exactly match the value tensor's shape before calling scatter (use `indices.expand(values.shape())` or repeat along mismatched dims).
2. Or reshape/slice the value tensor to the index tensor's shape.
3. If porting PyTorch code, note burn's scatter does not broadcast — construct both tensors with identical shapes explicitly.
4. Add a debug assert on `indices.shape() == values.shape()` in your calling code to fail earlier with your own message.

Example fix

// before
let out = x.scatter(1, &idx /* [N,1] */, &vals /* [N,K] */); // panics
// after
let idx_full = idx.expand(vals.shape());
let out = x.scatter(1, &idx_full, &vals);

fn ensure_scatter_shapes<B: Backend>(indices: &Tensor<B,2>, values: &Tensor<B,2>) -> Result<(), String> {
    (indices.shape() == values.shape())
        .then_some(())
        .ok_or_else(|| format!("scatter shape mismatch: indices {:?} vs values {:?}",
            indices.shape(), values.shape()))
}

let out = std::panic::catch_unwind(AssertUnwindSafe(|| x.scatter(dim, &idx, &vals)))
    .map_err(|_| anyhow!("scatter requires indices.shape() == values.shape()"))?;

When it happens

Trigger: Calling `scatter(dim, indices, values)` where `indices.dims() != values.dims()`; e.g. indices of shape [N, K] with values of shape [N, K'] or a squeezed/expanded index tensor.

Common situations: Hand-building scatter inputs where indices were gathered for a differently-shaped value set; off-by-one or wrong-dim expansion when broadcasting indices; porting from PyTorch's scatter (which broadcasts) into burn's stricter API.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• Expected float handle, got {}
• Expected int handle, got {}
• Expected bool handle, got {}
• Expected quantized handle, got {}
• ndarray scatter_nd requires contiguous values

AI-assisted analysis of tracel-ai/burn@d16f7ba2ed (2026-09-05). Data as JSON: /api/errors/371114581634fcf5. Report an issue: GitHub.