tracel-ai/burn · error

float_select_assign with {other:?} update is not implemented

Error message

float_select_assign with {other:?} update is not implemented

What it means

In the autodiff backend's float_select_assign (crates/burn-autodiff/src/ops/tensor.rs:1964), only Add/Sub/Mul IndexingUpdateOps have gradient implementations; any other update op hits the catch-all arm and panics with unimplemented!("float_select_assign with {other:?} update is not implemented").

Source

Thrown at crates/burn-autodiff/src/ops/tensor.rs:1964

                                tensor.primitive,
                                dim,
                                indices,
                                value.primitive,
                                IndexingUpdateOp::Mul,
                            ),
                        )
                    }
                    OpsKind::UnTracked(prep) => prep.finish(B::float_select_assign(
                        tensor.primitive,
                        dim,
                        indices,
                        value.primitive,
                        IndexingUpdateOp::Mul,
                    )),
                }
            }
            other => {
                unimplemented!("float_select_assign with {other:?} update is not implemented")
            }
        }
    }

    fn float_slice(tensor: FloatTensor<Self>, slices: &[Slice]) -> FloatTensor<Self> {
        #[derive(Debug)]
        struct Index;

        #[derive(new, Debug)]
        struct RetroSlice<B: Backend> {
            tensor_id: NodeId,
            slices: Vec<Slice>,
            _backend: PhantomData<B>,
        }

        impl<B: Backend> RetroForward for RetroSlice<B> {
            fn forward(&self, states: &mut BackwardStates, out_node: NodeId) {
                let tensor = states.get_state::<B::FloatTensorPrimitive>(&self.tensor_id);

View on GitHub (pinned to d16f7ba2ed)

Solutions

  1. Use Add/Sub/Mul select_assign updates in autodiff-tracked code.
  2. Detach the tensor (no gradient tracking) before performing the unsupported update.
  3. Implement the missing gradient rule for the update op in burn-autodiff.

Example fix

// before
embeddings.select_assign(0, indices, updates, IndexingUpdateOp::Assign);
// after
embeddings.select_assign(0, indices, updates, IndexingUpdateOp::Add);
Defensive patterns

Strategy: validation

Validate before calling

assert!(matches!(op, IndexingUpdateOp::Add | IndexingUpdateOp::Sub | IndexingUpdateOp::Mul), "autodiff select_assign only supports Add/Sub/Mul");

Prevention

When it happens

Trigger: Calling Tensor::select_assign (or float_select_assign) with an IndexingUpdateOp other than Add, Sub, or Mul under burn-autodiff.

Common situations: Assign-style index updates inside a training loop (e.g. updating embeddings by direct assignment); the same op works forward-only but panics when gradients are tracked.

Related errors


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