tracel-ai/burn · error

Unsupported dimension, only the last dimension can differ: T

Error message

Unsupported dimension, only the last dimension can differ: Tensor {:?} Index {:?}

What it means

gather (and scatter batch computation) allows only the last dimension of the index tensor to differ from the tensor's shape; all leading (batch) dimensions must match. gather_batch_size panics when any leading dimension of shape_tensor differs from the corresponding dimension of shape_indices.

Source

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

            let out_offset = n * slice_size;
            output_vec[out_offset..(out_offset + slice_size)]
                .copy_from_slice(&data_flat[base_offset..(base_offset + slice_size)]);
        }

        let out_shape = Shape::from(out_shape_vec);
        let output = ArrayD::from_shape_vec(out_shape.as_slice(), output_vec)
            .expect("gather_nd: shape mismatch");

        output.into_shared()
    }

    fn gather_batch_size(shape_tensor: &[usize], shape_indices: &[usize]) -> usize {
        let ndims = shape_tensor.num_dims();
        let mut batch_size = 1;

        for i in 0..ndims - 1 {
            if shape_tensor[i] != shape_indices[i] {
                panic!(
                    "Unsupported dimension, only the last dimension can differ: Tensor {:?} Index \
                     {:?}",
                    shape_tensor, shape_indices
                );
            }
            batch_size *= shape_indices[i];
        }

        batch_size
    }

    pub fn reshape(tensor: SharedArray<E>, shape: Shape) -> SharedArray<E> {
        reshape!(
            ty E,
            shape shape,
            array tensor,
            d shape.num_dims()
        )

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Make all leading dimensions of the indices tensor equal to the input tensor's leading dimensions.
2. Reshape or slice the indices so only the last dimension differs from the tensor shape.
3. Adjust the batch handling code so indices are produced per-batch with the same batch size.
4. Print tensor.dims() and indices.dims() and align them before calling gather.

Example fix

// before
let tensor = Tensor::zeros([2, 3, 4], &device);
let indices = Tensor::zeros([8, 3, 2], &device); // wrong leading dim
tensor.gather(2, indices);
// after
let indices = Tensor::zeros([2, 3, 2], &device); // leading dims match
tensor.gather(2, indices);

let (td, id) = (tensor.dims(), indices.dims());
assert_eq!(td.len(), id.len());
assert!(td[..td.len()-1] == id[..id.len()-1], "gather: leading dims must match");

fn gather_ok(t: &[usize], i: &[usize]) -> bool {
    t.len() == i.len() && t[..t.len()-1] == i[..i.len()-1]
}

When it happens

Trigger: Calling Tensor::gather with an indices tensor whose leading dimensions don't match the input tensor, e.g. tensor of shape [2, 3, 4] gathered with indices of shape [5, 3, 2].

Common situations: Batch size mismatch between data and indices (wrong batch slice, dropped/added batch dimension); passing 2D indices for a 3D tensor; porting numpy fancy-indexing code that supports arbitrary shapes.

Related errors

• Incompatible shapes for broadcasting: {:?} and {:?}
• gather_nd is not supported for bool tensors
• gather_nd: shape mismatch
• capture tensor operations must run inside CaptureDevice::cap
• capture tensor {} has no initialized value

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