tracel-ai/burn · error

argmax: unsupported dtype {:?}

Error message

argmax: unsupported dtype {:?}

What it means

The burn-flex `argmax` reduction panics when the tensor dtype has no argmax implementation. The match at reduce.rs:742-780 handles F32/F64/F16/BF16 and I8-I64 only; unsigned integers (U8/U16/U32/U64) and Bool fall into the `_` arm and panic with the dtype name. Note the asymmetry with `min`/`max`, which do support unsigned types.

Source

Thrown at crates/burn-flex/src/ops/reduce.rs:779

                f16::from_f32,
            )
            .1
        }
        DType::BF16 => {
            extremum_dim_with_indices_half::<bf16, _>(
                &tensor,
                dim,
                |a, b| !b.is_nan() && (a.is_nan() || a > b),
                bf16::to_f32,
                bf16::from_f32,
            )
            .1
        }
        DType::I8 => extremum_dim_with_indices::<i8, _>(&tensor, dim, |a, b| a > b).1,
        DType::I16 => extremum_dim_with_indices::<i16, _>(&tensor, dim, |a, b| a > b).1,
        DType::I32 => extremum_dim_with_indices::<i32, _>(&tensor, dim, |a, b| a > b).1,
        DType::I64 => extremum_dim_with_indices::<i64, _>(&tensor, dim, |a, b| a > b).1,
        _ => panic!("argmax: unsupported dtype {:?}", tensor.dtype()),
    }
}

/// Argmin along a dimension, returning indices as isize (INDEX_DTYPE).
pub fn argmin(tensor: FlexTensor, dim: usize) -> FlexTensor {
    assert!(
        tensor.layout().shape()[dim] > 0,
        "argmin: dimension {dim} has size 0"
    );
    assert_dim_fits_isize(tensor.layout().shape()[dim], dim);
    // f32 last-dim fast path: 2-pass SIMD for large rows, 1-pass scalar for small rows
    if tensor.dtype() == DType::F32 && dim == tensor.layout().shape().num_dims() - 1 {
        #[cfg(feature = "simd")]
        if tensor.layout().shape()[dim] >= EXTREMUM_SIMD_ROW_THRESHOLD {
            return extremum_indices_f32_last_simd(&tensor, dim, kernels::min_f32);
        }
        return extremum_indices_f32_last_scalar(&tensor, dim, |a, b| a < b);
    }

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast unsigned input to a signed or float type first: `tensor.cast(DType::I64)` or `tensor.cast(DType::F32)` before argmax (values must fit; u64 above i64::MAX overflows).
2. If the tensor is Bool, cast to U8 (or I64) first, then argmax.
3. If you need unsigned argmax, add U8/U16/U32/U64 arms to the match in crates/burn-flex/src/ops/reduce.rs using `extremum_dim_with_indices::<u32, _>(...)` etc.
4. Verify the producing op's output dtype; if you intended floats, fix the cast upstream instead of casting at the argmax call.

Example fix

// before
let idx = argmax(u8_tensor, 1); // panics: unsupported dtype U8
// after
let idx = argmax(u8_tensor.cast(DType::I64), 1);

fn ensure_argmax_supported(dtype: DType) -> Result<(), String> {
    match dtype {
        DType::F32 | DType::F64 | DType::F16 | DType::BF16
        | DType::I8 | DType::I16 | DType::I32 | DType::I64 => Ok(()),
        other => Err(format!("argmax: unsupported dtype {other:?}; cast first")),
    }
}

fn is_argmax_supported(dtype: DType) -> bool {
    matches!(dtype, DType::F32 | DType::F64 | DType::F16 | DType::BF16
        | DType::I8 | DType::I16 | DType::I32 | DType::I64)
}

let out = std::panic::catch_unwind(|| argmax(t.clone(), dim))
    .ok()
    .unwrap_or_else(|| argmax(t.cast(DType::I64), dim));

When it happens

Trigger: Calling `ops::reduce::argmax(tensor, dim)` on a U8/U16/U32/U64 or Bool tensor; also on any dtype not in the supported set. The dim bounds and size-0 cases are caught by earlier asserts, so this panic is purely dtype-driven.

Common situations: Argmax over an unsigned tensor loaded from a file/quantized pipeline (u8 images, u32 ids); argmax over a bool mask to find index of first/any true; switching backends where the other backend supported unsigned argmax but burn-flex does not.

Related errors

• min: unsupported dtype {:?}
• argmin: unsupported dtype {:?}
• Quantization scheme is not valid for dtype {other:?}
• Can't store native sub-byte values
• burn-flex does not support Bool(U32) storage (only Native an

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