tracel-ai/burn · error

Reshape would split a block across multiple rows.

Error message

Reshape would split a block across multiple rows.

What it means

For 1D block-quantized tensors, if the reshaped last dimension becomes smaller than the block size, one quantization block would span multiple rows, which cannot be represented by per-row scales unless there is only a single block overall. When scales hold more than one block, the op panics.

Source

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

                    unimplemented!(
                        "Split reshape of ND block-quantized tensor is not yet supported."
                    );
                }
            }
            other => unreachable!("Reshape analysis {other:?} should not update strides."),
        }
    }

    let shape_last = *shape.last().unwrap();

    // The per-tensor scale is a scalar in its own region, so only the block grid moves.
    let shape_scales = match scheme.block_size() {
        None => scales.meta.shape().clone(), // always [1], invariant under reshape
        Some(block_size) if block_size.len() == 1 && shape_last < (block_size[0] as usize) => {
            // If the new last dimension is smaller than the block size,
            // it means a single block now spans across multiple rows.
            if scales.meta.shape().num_elements() > 1 {
                unimplemented!("Reshape would split a block across multiple rows.");
            }
            // Exception: allow if there is exactly 1 block total (essentially per-tensor quantization)
            scales.meta.shape().clone()
        }
        Some(_) => {
            // ND blocks: derive scales shape from the new tensor shape
            params_shape(&shape, &scheme)
        }
    };

    let action_scales = reshape_action(scales.meta.shape(), scales.meta.strides(), &shape_scales);

    match (action_values, action_scales) {
        (
            ReshapeAction::UpdateStrides { strides },
            ReshapeAction::UpdateStrides {
                strides: scales_strides,
            },

View on GitHub (pinned to d16f7ba2ed)

Solutions

  1. Keep the last dimension a multiple of (or >=) the block size when reshaping
  2. Dequantize, reshape, and re-quantize with a suitable block size
  3. Re-quantize with a smaller block size matching the new last dimension
  4. Use per-tensor quantization if heavy reshaping is required

Example fix

// before
let y = x.reshape([4, 32]); // block_size=64, scales>1 -> panic
// after
let y = x.dequantize().reshape([4, 32]); // or keep dim >= 64
Defensive patterns

Strategy: validation

Validate before calling

fn last_dim_keeps_block(new_last: usize, block_size: Option<&[usize]>) -> bool {
    block_size.map_or(true, |b| new_last >= b[0] as usize || new_last == 0)
}

Type guard

fn block_fits_in_last_dim(new_shape: &[usize], scheme: &QuantScheme) -> bool {
    scheme.block_size().map_or(true, |b| {
        *new_shape.last().unwrap() >= b[0] as usize
    })
}

Try / catch

// Guard:
if block_fits_in_last_dim(&new_shape, &scheme) { x.reshape(new_shape) } else { x.dequantize().reshape(new_shape) }

Prevention

When it happens

Trigger: Reshaping a per-block (1D block_size, e.g. 64) quantized tensor so its last dimension shrinks below the block size (e.g. [4,256] -> [4,32] with block 64) while scales contain multiple blocks.

Common situations: Downsizing the feature dimension of block-quantized weights/activations; folding dimensions in custom layers that bypass quantization-aware shape planning.

Related errors


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