Dimension mismatch: cannot broadcast dimension {tensor_dim}

Error message

Dimension mismatch: cannot broadcast dimension {tensor_dim} of tensor to target shape

What it means

expand creates a broadcast view of a tensor toward a target shape; each tensor dimension must either equal the target dimension or be 1 (in which case stride 0 is used). If a tensor dimension is neither 1 nor equal to the target, the stride computation cannot proceed and the library panics. This is a shape-contract violation detected client-side before any kernel launch.

Source

Thrown at crates/burn-cubecl/src/ops/base.rs:275

    // Calculate the difference in dimensions
    let dim_diff = ndims_out.saturating_sub(ndims_in);

    // Compare dimensions from the end, setting strides for matching dimensions or broadcasted ones
    let mut tensor_dim_iter = tensor.meta.shape().iter().rev();
    for i in (0..ndims_out).rev() {
        if i >= dim_diff {
            if let Some(&tensor_dim) = tensor_dim_iter.next() {
                if tensor_dim == target_shape[i] || tensor_dim == 1 {
                    // Copy stride for non-broadcast dimensions or set to 0 for broadcast ones
                    new_strides[i] = if tensor_dim == target_shape[i] {
                        tensor.meta.strides()[i - dim_diff]
                    } else {
                        0
                    };
                } else {
                    // Error handling: Dimension mismatch for broadcasting
                    panic!(
                        "Dimension mismatch: cannot broadcast dimension {tensor_dim} of tensor to target shape"
                    );
                }
            } else {
                // If the input tensor has fewer dimensions, treat missing dimensions as 1
                // and set stride to 0 (broadcasting)
                new_strides[i] = 0;
            }
        } else {
            // For extra dimensions in the target shape, set stride to 0 (broadcasting)
            new_strides[i] = 0;
        }
    }

    // Extra check to ensure block scales must be properly handled once they're added
    if tensor.qparams.is_some() && tensor.scheme().block_size().is_some() {
        todo!()
    }

View on GitHub (pinned to d16f7ba2ed)

Solutions

Make each tensor dim either 1 or equal to the target dim before expanding
Insert an unsqueeze first so singleton dims line up with the target shape's trailing dims
Print both shapes and compare right-aligned (broadcasting aligns trailing dims)
Compute the target shape from the input at runtime instead of hard-coding it

Example fix

// before: expand([3,4] -> [8,4]) panics
let y = x.expand([8, 4]);
// after: expand along a dim that is 1
let x = Tensor::ones([1, 4]);
let y = x.expand([8, 4]);

Defensive patterns

Strategy: validation

Validate before calling

fn can_expand(shape: &[usize], target: &[usize]) -> bool {
    let off = target.len() - shape.len();
    shape.iter().zip(&target[off..]).all(|(s, t)| *s == 1 || s == t)
}

Prevention

Always align shapes right-aligned mentally when broadcasting
Insert unsqueeze for missing leading dims before expand
Derive expand targets from input shapes at runtime

When it happens

Trigger: Calling bool_expand/int_expand/float_expand (or expand via them) with a target shape whose dimension i is > 1 and != tensor's dimension i, for some aligned trailing dimensions.

Common situations: Expanding [3, 4] to [8, 4], mixing batch dims like expanding [B, 1] to [B2, C] with wrong ordering, or hard-coded target shapes that don't match runtime batch sizes.

Related errors

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