What causes "gather_nd is not supported for bool tensors" in tracel-ai/burn?

Cast the bool tensor to a numeric dtype (e.g. u8/f32) before gather_nd, then convert back.

gather_nd is not supported for bool tensors

Error message

gather_nd is not supported for bool tensors

What it means

Capability guard on the bool tensor bridge: `gather_nd` (like `scatter_nd`) is deliberately unsupported for BoolKind bridge tensors — bool tensors cannot be gathered with an ND index op in this backend abstraction — so any call panics. The failing input is a bool tensor passed to gather_nd where a numeric dtype was expected.

Source

Thrown at crates/burn-tensor/src/bridge/ops/bool.rs:167

                tensor.into(),
                indices.into(),
                values.into(),
            )),
            _ => unimplemented!(),
        }
    }

    fn scatter_nd(
        _data: BridgeTensor,
        _indices: BridgeTensor,
        _values: BridgeTensor,
        _reduction: IndexingUpdateOp,
    ) -> BridgeTensor {
        panic!("scatter_nd is not supported for bool tensors")
    }

    fn gather_nd(_data: BridgeTensor, _indices: BridgeTensor) -> BridgeTensor {
        panic!("gather_nd is not supported for bool tensors")
    }

    fn device(tensor: &BridgeTensor) -> Device {
        Device::new(tensor.as_dispatch().device())
    }

    fn to_device(tensor: BridgeTensor, device: &Device) -> BridgeTensor {
        BridgeTensor::bool(Dispatch::bool_to_device(
            tensor.into(),
            device.as_dispatch(),
        ))
    }

    async fn into_data_async(tensor: BridgeTensor) -> Result<TensorData, ExecutionError> {
        Dispatch::bool_into_data(tensor.into()).await
    }

    fn from_data(data: TensorData, device: &Device, dtype: DType) -> BridgeTensor {

View on GitHub (pinned to d16f7ba2ed)

Solutions

Cast the bool tensor to a numeric dtype (e.g. u8/f32) before gather_nd, then convert back.
Perform the gather manually with boolean masking or index_select, which are supported for bools.
Restructure the operation so indices are applied to the numeric source tensor before the bool conversion.

Example fix

// before
let picked: Tensor<B, 2, Bool> = mask.gather(indices); // gather_nd on Bool
// after
let picked = mask.int().gather(indices).bool_cast();

Defensive patterns

Strategy: fallback

Validate before calling

fn gather_bool_mask<B: Backend, const D: usize>(mask: Tensor<B, D, Bool>, indices: Tensor<B, 2, Int>) -> Tensor<B, D, Bool> {
    mask.int().gather(indices).bool_cast() // gather_nd path via Int fallback
}

Try / catch

std::panic::catch_unwind(|| mask.gather(indices))
    .map_err(|_| "gather_nd unsupported for Bool")
    .or_else(|_| Ok::<_, String>(mask.int().gather(indices).bool_cast()))?;

Prevention

Cast Bool tensors to Int before gather_nd-style operations and cast back afterwards
Prefer select/index along dims, which Bool supports, when gathering masks
Flag ONNX GatherND-on-bool nodes during model import and rewrite them to int gathers

When it happens

Trigger: Calling tensor.gather(indices) / gather_nd on a Tensor::<B, D, Bool> — e.g. pulling selected boolean mask entries by index tensors.

Common situations: Reusing numeric gather code paths on boolean masks; model translation (ONNX GatherND on bool) into burn with the Bool element type.

Related errors

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