What causes "{other:?} reduction is not supported" in tracel-ai/burn?

Use MseLoss::forward_no_reduction(logits, targets) to get the unreduced tensor and call .mean()/.sum() or apply custom weighting yourself.

tracel-ai/burn · error

{other:?} reduction is not supported

Error message

{other:?} reduction is not supported

What it means

MseLoss::forward supports only Mean, Auto, and Sum reductions; any other Reduction variant panics. MseLoss computes mean squared error element-wise via forward_no_reduction then applies the reduction, so this panic comes from the catch-all arm of the match.

Source

Thrown at crates/burn-nn/src/loss/mse.rs:40

    }

    /// Compute the criterion on the input tensor.
    ///
    /// # Shapes
    ///
    /// - logits: [batch_size, num_targets]
    /// - targets: [batch_size, num_targets]
    pub fn forward<const D: usize>(
        &self,
        logits: Tensor<D>,
        targets: Tensor<D>,
        reduction: Reduction,
    ) -> Tensor<1> {
        let tensor = self.forward_no_reduction(logits, targets);
        match reduction {
            Reduction::Mean | Reduction::Auto => tensor.mean(),
            Reduction::Sum => tensor.sum(),
            other => panic!("{other:?} reduction is not supported"),
        }
    }

    /// Compute the criterion on the input tensor without reducing.
    pub fn forward_no_reduction<const D: usize>(
        &self,
        logits: Tensor<D>,
        targets: Tensor<D>,
    ) -> Tensor<D> {
        logits.sub(targets).square()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use burn::tensor::TensorData;

View on GitHub (pinned to d16f7ba2ed)

Solutions

Use MseLoss::forward_no_reduction(logits, targets) to get the unreduced tensor and call .mean()/.sum() or apply custom weighting yourself.
Pass Reduction::Mean, Reduction::Sum, or Reduction::Auto to forward.
Audit configs/tests for Reduction::None and replace with explicit no-reduction APIs.

Example fix

// before
let loss = mse.forward(predictions, targets, Reduction::None); // panics
// after
let per_elem = mse.forward_no_reduction(predictions, targets);
let loss = per_elem.mean(); // or custom per-sample reduction

Defensive patterns

Strategy: validation

Validate before calling

fn ensure_supported(r: &Reduction) -> Result<(), String> {
    match r {
        Reduction::Mean | Reduction::Auto | Reduction::Sum => Ok(()),
        other => Err(format!("unsupported reduction for MSE: {other:?}")),
    }
}

Type guard

fn is_supported_reduction(r: &Reduction) -> bool {
    matches!(r, Reduction::Mean | Reduction::Auto | Reduction::Sum)
}

Try / catch

let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(||
    mse.forward(predictions, targets, reduction.clone())));

Prevention

Never pass Reduction::None to MseLoss::forward; use forward_no_reduction instead.
For per-sample or mask-weighted MSE, compute it from forward_no_reduction output.
Centralize reduction validation for all losses in one config check.
Update tests (e.g. test_mse_loss) to only use supported variants.

When it happens

Trigger: Calling MseLoss::forward(logits, targets, Reduction::None) — e.g. from test_mse_loss or training code expecting per-element MSE — or any reduction value other than Mean/Auto/Sum.

Common situations: Porting PyTorch MSELoss(reduction='none') for per-sample or mask-weighted losses; sharing one Reduction config across losses; upgrading burn and finding Reduction::None no longer handled in forward.

Related errors

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