Implement remaining IndexingUpdateOp variants for scatter/select_assign across backends

[antimora]

Feature description

Follow-up to #4477. IndexingUpdateOp now has all five variants (Assign, Add, Mul, Min, Max) and both Tensor::scatter and Tensor::select_assign take it, but most backends still only implement a subset and fall through to unimplemented!() at runtime.

scatter_nd / select-family aside, here is the current state on main for the element-wise ops:

scatter (element-wise)

Backend	Assign	Add	Mul	Min	Max
ndarray	yes	yes	yes	no	no
flex	no	yes	yes	no	no
cubecl (wgpu/cuda/metal/rocm)	no	yes	yes	no	no
tch	no	yes	yes	no	no
autodiff	yes	yes	yes	no	no
bool (bridge)	no	yes (or)	no	no	no

select_assign

Backend	Assign	Add	Mul	Min	Max
ndarray	yes	yes	yes	no	no
flex	no	yes	yes	no	no
cubecl	no	yes	yes	no	no
tch	no	yes	yes	no	no
bool (bridge)	no	yes (or)	no	no	no

scatter_nd (for contrast, already complete)

All five reductions are implemented in ndarray, flex and cubecl. This is the shape the element-wise ops should reach.

So the concrete gaps are:

• Assign missing in burn-flex, burn-cubecl and burn-tch for both scatter and select_assign (ndarray has it).
• Min / Max missing in every backend for both scatter and select_assign.

The default trait bodies in burn-backend/src/backend/ops/tensor.rs (float_scatter, float_select_assign) forward only Add, so a backend that does not override them silently inherits Add-only support and panics at runtime rather than failing to compile.

Feature motivation

This is what blocks the ONNX side from finishing tracel-ai/burn-onnx#486.

ONNX ScatterElements has five reduction modes that map one-to-one onto IndexingUpdateOp. ScatterND can already lower all five to scatter_nd because that op is complete. ScatterElements cannot:

ONNX reduction	Lowering today in burn-onnx
add	scatter(..., Add)
mul	scatter(..., Mul)
none	gather-diff-scatter-add workaround, because Assign is unavailable on the default CPU backend (flex)
max	data + (updates - gathered).clamp_min(0) scattered with Add
min	data + (updates - gathered).clamp_max(0) scattered with Add

The none/max/min rewrites are correct for unique indices but diverge from ONNX semantics when indices contains duplicates, where ONNX defines the reduction as applied repeatedly. They also do an extra gather pass over the data. Native Assign/Min/Max would remove both problems and let all five modes be a single scatter() call.

Bool tensors are a smaller related gap: the bridge only maps Add to bool_scatter_or, and or cannot clear a bit, so ONNX ScatterElements/ScatterND on bool has to round-trip through i64.

Suggest a Solution

• Implement Assign for scatter and select_assign in burn-flex, burn-cubecl and burn-tch.
• Implement Min / Max for scatter and select_assign across the backends. The cubecl side already has BinaryMinOp / BinaryMaxOp wired up for scatter_nd_kernel, so the element-wise scatter_kernel should be able to reuse them.
• Extend burn-autodiff to cover whichever of these are differentiable (Min/Max gradients route to the argmin/argmax side, same as scatter_nd).
• Consider making the missing combinations a compile-time or documented capability rather than a runtime unimplemented!(), so downstream code generators can pick a supported lowering instead of discovering the gap on the first forward pass.

[antimora]

Update on the ONNX side: this is no longer a correctness blocker, only a performance one.

ScatterElements in burn-onnx now lowers none/mul/min/max to scatter_nd instead of element-wise scatter, since scatter_nd already implements all five update ops on every backend. ScatterElements writes to data[p_0, .., indices[p], .., p_{r-1}] for every p in the index shape, so materializing those coordinates as index tuples turns it into a ScatterND.

That works today, but it pays for building an [n, r] coordinate tensor (a flat arange decomposed by row-major strides, one column per rank) on every call, purely to work around the missing element-wise ops. Once scatter gains Assign/Mul/Min/Max, all four collapse back to a single scatter(dim, indices, updates, op) call with no index construction at all.

The gaps in the tables above are unchanged and still accurate. Correcting one thing in my original motivation section: the delta-based rewrites I described there (data + (updates - gathered) scattered with Add) are not a viable workaround and are gone. e + (max(e,u) - e) is NaN when e = -inf, the intermediate subtraction overflows narrow integer types, and duplicate indices sum their deltas instead of folding with the reduction.

Ref: tracel-ai/burn-onnx#487

[ax1s-x1zz]

Small update on the current state of main: the element-wise scatter / select_assign Assign gap is no longer there — burn-flex, burn-cubecl and burn-tch have all been extended since this issue was opened (ndarray already had it). What remains is exactly the Min / Max row:

• Min / Max are only implemented for scatter_nd; every element-wise scatter / select_assign (float and int) with Min / Max hits unimplemented! in all four backends.
• burn-autodiff has Min/Max backward only for scatter_nd, not for the element-wise ops.

I'm working on implementing Min / Max for the element-wise scatter and select_assign in the ndarray, flex, cubecl and tch backends plus the autodiff backward, mirroring the existing scatter_nd implementations.