Trying to consume the gradients for an untracked tensor

Error message

Trying to consume the gradients for an untracked tensor

What it means

Gradients container lookup failure: `consume` only retrieves gradients for nodes whose autodiff `Requirement` is `Grad`. Calling `consume` on a tensor whose gradients were never registered (created with grad tracking disabled, or already consumed/removed from the container) means the `.get()` returns None and the expect fires — a programming error in backward-pass bookkeeping, not a recoverable runtime condition.

Source

Thrown at crates/burn-autodiff/src/grads.rs:108

    }

    /// Consumes the gradients for a given tensor.
    ///
    /// Each tensor should be consumed exactly 1 time if its gradients are only required during the
    /// backward pass, otherwise, it may be consume multiple times.
    pub fn consume<B: Backend>(&mut self, node: &NodeRef) -> FloatTensor<B> {
        match node.requirement {
            Requirement::Grad => self
                .container
                .get::<TensorPrimitive<B>>(&node.id.value)
                .map(|tensor| tensor.tensor())
                .expect("Can't consume the gradients before they are registered at least once."),
            Requirement::GradInBackward => self
                .container
                .remove::<TensorPrimitive<B>>(&node.id.value)
                .map(|tensor| tensor.tensor())
                .expect("Can't consume the gradients before they are registered at least once."),
            Requirement::None => panic!("Trying to consume the gradients for an untracked tensor"),
        }
    }

    /// Removes a grad tensor from the container.
    pub fn remove<B: Backend>(&mut self, tensor: &AutodiffTensor<B>) -> Option<FloatTensor<B>> {
        self.container
            .remove::<TensorPrimitive<B>>(&tensor.node.id.value)
            .map(|tensor| tensor.tensor())
    }

    /// Gets a grad tensor from the container.
    pub fn get<B: Backend>(&self, tensor: &AutodiffTensor<B>) -> Option<FloatTensor<B>> {
        self.container
            .get::<TensorPrimitive<B>>(&tensor.node.id.value)
            .map(|tensor| tensor.tensor())
    }

    /// Register a grad tensor in the container.

View on GitHub (pinned to d16f7ba2ed)

Solutions

Switch the update op to IndexingUpdateOp::Add so the bool_select_or path is used.
Do the select-update on an int tensor and cast to bool.
Extend burn-router's bool select match arm if a new op needs support upstream.

Example fix

// before
let out = tensor_bool.select(dim, indices, value, IndexingUpdateOp::Set); // panics
// after
let out = tensor_bool.select(dim, indices, value, IndexingUpdateOp::Add);

Defensive patterns

Strategy: validation

Validate before calling

fn bool_select_supported(update: IndexingUpdateOp) -> bool {
    matches!(update, IndexingUpdateOp::Add)
}
// assert!(bool_select_supported(desc.update)) before select on a bool tensor

Type guard

fn is_add_update(op: &IndexingUpdateOp) -> bool {
    matches!(op, IndexingUpdateOp::Add)
}

Try / catch

std::panic::catch_unwind(std::panic::AssertUnwindSafe(||
    run_select_on_bool(...)
)).map_err(|_| anyhow::anyhow!("bool select only supports Add update op"))

Prevention

Use Add (OR) update semantics for bool select.
Compute selects on int tensors, convert to bool afterward.
Keep traced graphs limited to ops the router implements for bool.
Add graph linting to flag bool select ops with non-Add updates.

When it happens

Trigger: Executing a Select operation IR on a bool tensor whose desc.update is not IndexingUpdateOp::Add (e.g. a select/index assignment with a non-additive update op traced into the graph).

Common situations: Models using gather/select-and-update patterns on bool tensors routed through burn-router; code generated from frameworks where bool select defaults to a non-Add update op.

Related errors

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