tracel-ai/burn · error

softmax: unsupported dtype {:?}

Error message

softmax: unsupported dtype {:?}

What it means

burn-flex's softmax_last dispatches on the tensor's DType and has kernels only for F32, F64, F16 and BF16; any other dtype (integer, bool, etc.) hits the catch-all arm and panics. The library intentionally fails fast instead of silently casting, since softmax is only defined for floating-point data. This is a hard panic, not a Result, so it aborts the calling thread immediately.

Source

Thrown at crates/burn-flex/src/ops/activation.rs:229

    );

    if dim != rank - 1 {
        let swapped = Flex::float_swap_dims(tensor, dim, rank - 1);
        let normed = softmax_last(swapped);
        return Flex::float_swap_dims(normed, dim, rank - 1);
    }

    softmax_last(tensor)
}

fn softmax_last(tensor: FloatTensor<Flex>) -> FloatTensor<Flex> {
    let tensor = tensor.to_contiguous();
    match tensor.dtype() {
        DType::F32 => softmax_last_f32(tensor),
        DType::F64 => softmax_last_f64(tensor),
        DType::F16 => softmax_last_f16(tensor),
        DType::BF16 => softmax_last_bf16(tensor),
        dtype => panic!("softmax: unsupported dtype {:?}", dtype),
    }
}

fn softmax_last_f32(tensor: FlexTensor) -> FlexTensor {
    let shape = tensor.layout().shape().clone();
    let last = *shape.last().expect("softmax: empty shape");
    if last == 0 {
        return tensor;
    }
    let input: &[f32] = tensor.storage();
    let n = input.len();

    // Zero-initialize the output. The previous implementation used
    // `Vec::with_capacity` + `spare_capacity_mut` + a raw-pointer cast to
    // `&mut [f32]` to skip the memset, but forming a `&mut [f32]` over
    // uninitialized memory violates Rust's validity invariant (references
    // must point to initialized values of the correct type) even if every
    // element is written before it is read. The sound zero-memset

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast the tensor to a float dtype before softmax, e.g. tensor.cast(DType::F32) (or .float() on the burn tensor API).
2. Check where the input was created and fix the producing op so it emits F32/F64/F16/BF16 logits directly.
3. If you control the dispatch, extend softmax_last with a cast-to-f32 wrapper for unsupported dtypes instead of panicking.

Example fix

// before
let probs = tensor.softmax(); // tensor is I64
// after
let probs = tensor.cast(burn_std::DType::F32).softmax();

if !matches!(tensor.dtype(), DType::F32 | DType::F64 | DType::F16 | DType::BF16) {
    tensor = tensor.cast(DType::F32);
}
let probs = softmax(tensor);

fn is_float_dtype(d: DType) -> bool {
    matches!(d, DType::F32 | DType::F64 | DType::F16 | DType::BF16)
}

// it panics (not a Result); catch only at a recovery boundary
let probs = std::panic::catch_unwind(|| softmax(tensor.clone()))
    .unwrap_or_else(|_| softmax(tensor.cast(DType::F32)));

When it happens

Trigger: Calling softmax (or backend softmax_last) with a FlexTensor whose dtype is an integer type (I64/I32/I16/I8/U64/U32/U16/U8) or Bool. Typically happens after an argmax-free pipeline where logits were produced by an integer cast, or a quantized model whose output tensor was never cast back to float.

Common situations: Quantized/inference pipelines forgetting to dequantize logits; a cast chain like .int() left in before normalization; mixing backends where an upstream op produced integer output; test fixtures building integer tensors and reusing them for softmax.

Related errors

• burn_flex::layer_norm: unsupported dtype {:?}
• attention: unsupported dtype {:?}
• binary_op: unsupported dtype {:?}
• scalar_op: unsupported dtype {:?}
• int_binary_op: unsupported dtype {:?}

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