tracel-ai/burn · error

Node {:?} is needed but never checkpointed

Error message

Node {:?} is needed but never checkpointed

What it means

In the router interpreter, bool scatter operations only support IndexingUpdateOp::Add (implemented via bool_scatter_or). Any other update op on a bool tensor hits unimplemented!(), because for booleans only the logical-OR accumulation is defined.

Source

Thrown at crates/burn-autodiff/src/checkpoint/builder.rs:225

        // from which we use the ids to find them again in the checkpointing actions
        for (node_id, n_required) in n_required_map {
            // We find the checkpointing action for node_id. It's likely in checkpointing_actions
            // so we check there first, otherwise it will be in backup.
            // Technically it can be there several times but can never be of both types, so we can assume the first we find is fine

            let action = match self
                .explicit_actions
                .iter()
                .position(|action| action.id() == node_id)
            {
                Some(pos) => self.explicit_actions.remove(pos),
                None => {
                    let pos = self
                        .backup_actions
                        .iter()
                        .position(|action| action.id() == node_id);
                    self.backup_actions.remove(pos.unwrap_or_else(|| {
                        panic!("Node {:?} is needed but never checkpointed", node_id)
                    }))
                }
            };

            match action {
                CheckpointingAction::Computed {
                    node_id: _,
                    state_content,
                } => {
                    self.checkpoint_compute(backward_states_map, node_id, state_content, n_required)
                }
                CheckpointingAction::Recompute {
                    node_id: _,
                    retro_forward,
                } => self.checkpoint_lazy(
                    backward_states_map,
                    retro_forward_map,
                    node_id,

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Use IndexingUpdateOp::Add semantics for bool scatter, or implement the update manually: invert the mask and multiply/mask the tensor.
2. Perform the scatter on an int/float tensor instead and convert back to bool afterward.
3. Only use bool_scatter_or-supported ops (index_put_ with Add) in traced code paths routed through burn-router.

Example fix

// before
tensor_bool.scatter(dim, indices, value, IndexingUpdateOp::Set); // panics
// after
tensor_bool.scatter(dim, indices, value, IndexingUpdateOp::Add); // bool_scatter_or path

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

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

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

When it happens

Trigger: Executing a Scatter operation IR on a bool tensor where desc.update is not IndexingUpdateOp::Add (e.g. it is a different IndexingUpdateOp variant such as a set/replace update coming from a traced scalar-op slice-assign).

Common situations: Tracing/exporting a model that performs tensor[index] = value (non-additive scatter) on bool tensors and running it through the burn-router; ops that are fine on float/int tensors but have no bool equivalent.

Related errors

• Trying to consume the gradients for an untracked tensor
• ctc_loss_backward: 2 * max_target_len + 1 = {} exceeds the k
• Invalid broadcast shapes: Next grad shape {:?}, Previous gra
• Can't differentiate embedding backward.
• Can't differentiate linear_x_backward.

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