tracel-ai/burn · error

gather_nd: shape mismatch

Error message

gather_nd: shape mismatch

What it means

After collecting all slices, gather_nd rebuilds the output ArrayD via from_shape_vec with the computed out_shape (idx_shape[..m-1] + data_shape[k..]). If the accumulated output_vec length does not equal num_indices * slice_size, from_shape_vec returns Err and this expect() panics - i.e. the computed output shape is inconsistent with the number of elements produced.

Source

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

        };

        let mut output_vec: Vec<E> = vec![0.elem::<E>(); out_total];

        for n in 0..num_indices {
            let mut base_offset = 0usize;
            for j in 0..k {
                let idx_val = idx_flat[n * k + j].elem::<i64>() as usize;
                base_offset += idx_val * strides[j];
            }

            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];
        }

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Verify indices rank: indices.dims().last() must be <= data.dims().len(); expected output shape = indices.dims()[..-1] + data.dims()[k..]
2. Validate all index values are within bounds of the corresponding data dimension before calling
3. Test with small known shapes to confirm k (slice depth) semantics; slice_size = product of data_shape[k..], num_indices = product of idx_shape[..m-1]

Example fix

// before
let idx = Tensor::<Cpu, 3>::from_data(...); // last dim 3, data rank 2
let out = data.gather(idx); // k > data rank -> inconsistent shape
// after
let idx = Tensor::<Cpu, 2>::from_data(...); // last dim <= data rank
let out = data.gather(idx);

// before calling gather, verify shape contract:
fn check_gather_shapes(data_dims: &[usize], idx_dims: &[usize]) -> Result<Vec<usize>, String> {
    let k = *idx_dims.last().ok_or("indices must be non-empty rank")?;
    if k > data_dims.len() { return Err(format!("indices last dim {} exceeds data rank {}", k, data_dims.len())); }
    let slice_size: usize = data_dims[k..].iter().product();
    let num_indices: usize = idx_dims[..idx_dims.len()-1].iter().product();
    Ok(idx_dims[..idx_dims.len()-1].iter().cloned().chain(data_dims[k..].iter().cloned()).collect())
    // expected output has num_indices * slice_size elements
}

When it happens

Trigger: Calling gather_nd where rank/index-dim parameters (k slices from indices last dim, m leading index dims) produce an out_total that doesn't match the filled output_vec; typically caused by indices whose shape doesn't match data's rank, or index values out of range writing past the intended region (indices out of bounds would first panic elsewhere; here the mismatch is shape arithmetic, e.g. empty index tensor edge cases or k/m mis-derivation).

Common situations: ONNX GatherND graphs with unusual index ranks; calling gather with indices whose last dimension exceeds data rank; zero-size index tensors where shape computation yields 0 elements but non-zero shape.

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

• ndarray gather_nd requires contiguous data
• ndarray gather_nd requires contiguous indices
• expand: cannot expand dimension {} from {} to {}
• float_gather: unsupported dtype {:?}
• Invalid dimension: the shape of the index tensor should be t

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