tracel-ai/burn · error

max_pool3d_backward: unsupported index dtype {other:?}

Error message

max_pool3d_backward: unsupported index dtype {other:?}

What it means

burn-flex's generated max_pool3d_backward dispatcher matches the pooling `indices` tensor dtype against signed integer types (I64, I32, I16, I8) and panics for anything else. Max-pool backward needs integer argmax indices; an unsigned, float, bool, or quantized index dtype has no kernel, so the autodiff backward pass aborts.

Source

Thrown at crates/burn-flex/src/ops/pool.rs:90

/// Generates adaptive_avg_pool3d typed dispatchers.
macro_rules! adaptive_avg_pool3d_typed {
    ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $div_fn:expr) => {
        pub fn $fn_name(x: FlexTensor, output_size: [usize; 3]) -> FlexTensor {
            adaptive_avg_pool3d_impl::<$T, _>(x, output_size, $dtype, $zero, $div_fn)
        }
    };
}

/// Generates max_pool3d_backward typed dispatchers.
macro_rules! max_pool3d_backward_typed {
    ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr) => {
        pub fn $fn_name(x: FlexTensor, grad: FlexTensor, indices: FlexTensor) -> FlexTensor {
            match indices.dtype() {
                DType::I64 => max_pool3d_backward_impl::<$T, i64>(x, grad, indices, $dtype, $zero),
                DType::I32 => max_pool3d_backward_impl::<$T, i32>(x, grad, indices, $dtype, $zero),
                DType::I16 => max_pool3d_backward_impl::<$T, i16>(x, grad, indices, $dtype, $zero),
                DType::I8 => max_pool3d_backward_impl::<$T, i8>(x, grad, indices, $dtype, $zero),
                other => panic!("max_pool3d_backward: unsupported index dtype {other:?}",),
            }
        }
    };
}

/// Generates avg_pool3d_backward typed dispatchers.
macro_rules! avg_pool3d_backward_typed {
    ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $div_fn:expr) => {
        pub fn $fn_name(
            x: FlexTensor,
            grad: FlexTensor,
            kernel_size: [usize; 3],
            stride: [usize; 3],
            padding: [usize; 3],
            count_include_pad: bool,
        ) -> FlexTensor {
            avg_pool3d_backward_impl::<$T>(
                x,

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Ensure the indices tensor passed to max_pool3d backward has a signed int dtype (I32 is the typical default): `indices.cast(DType::I32)`.
2. Check any serialization/checkpoint path for casts of indices to unsigned types and keep them signed.
3. If you control the forward call, leave indices in the backend's default int dtype instead of converting.

Example fix

// before
let indices = raw_indices.cast(DType::U32);
let grad_x = max_pool3d_backward(x, grad, indices);
// after
let indices = raw_indices.cast(DType::I32);
let grad_x = max_pool3d_backward(x, grad, indices);

assert!(matches!(indices.dtype(), DType::I64 | DType::I32 | DType::I16 | DType::I8), "max_pool3d_backward indices must be signed ints, got {:?}", indices.dtype());

fn is_signed_int(d: DType) -> bool {
    matches!(d, DType::I64 | DType::I32 | DType::I16 | DType::I8)
}

When it happens

Trigger: Running backward through `max_pool3d` when the saved indices tensor has a dtype outside I8–I64 — e.g. indices stored/cast as U32 or U8, or a float tensor mistakenly passed as indices.

Common situations: Custom checkpointing in burn-autodiff that serializes indices as unsigned; casting indices to unsigned for compact storage and forgetting to restore; hand-constructed grads feeding pool backward directly.

Related errors

• Autodiff float tensor is on the wrong backend (expected {bac
• Expected autodiff-wrapped float tensor for backend {backend}
• Autodiff should not wrap an autodiff device.
• an autodiff float primitive must have an enabled autodiff co
• autodiff context requires the `autodiff` feature

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