Contribution Proposal: ScatterND Implementation for PyTorch

🚀 The feature

Hey everyone, I’m Ali, a software engineer from Nuro, and we’ve been working on migrating a lot of vision models from tensorflow to pytorch. One particularly critical operation for these vision models is tf.tensor_scatter_nd_update , and its reduction counterparts (e.g. tf.tensor_scatter_nd_min, tf.tensor_scatter_nd_max , etc).

We'd like to offer to implement a robust implementation of ScatterND in pytorch. Some details:

• Uses only native torch layers, and is fully statically shaped
• Supports batching with torch.vmap if desired
• Really performant when used with torch.compile & torch-xla

I've attached sample implementations in the 'Additional context section'.

Motivation, pitch

The ScatterND operation was a major roadblock for the tensorflow => pytorch migration for Nuro. If this operation is widely available, more organizations might be encouraged to adopt pytorch (rather than JAX, which is the default migration path from tensorflow).

Alternatives

There are a few scatter_nd implementations scattered (pardon the pun) around the web but none seem to be perfect.

https://github.com/rusty1s/pytorch_scatter, this one uses custom CUDA kernels which isn't easy to use or maintain
https://gist.github.com/airalcorn2/c7846d6fcb58a30b25ea6d97e16fe025, this implementation has a bug
https://gist.github.com/Ending2015a/b034ebbedc55fec1d8ec3b7230a95f1e, doesn't support batching, or have static shapes which is important for torch.compile/torch-xla

Additional context

Forum post: https://discuss.pytorch.org/t/contribution-proposal-scatternd-implementation-for-pytorch/211874/1

(sample implementation) https://gist.github.com/boubezariali/2e3d9650461f302a541235e33d7cded2
(simple test) https://gist.github.com/boubezariali/a4a22736a14414404b45b17f62da9c2d

We also have extensive testing implemented that we're willing to contribute.

Start by reviewing the linked sample implementation and simple test in the Additional context section, then read the referenced forum proposal and compare the stated requirements for batching, static shapes, torch.compile, and torch-xla. Confirm the intended PyTorch API and scope with maintainers. Done means an agreed ScatterND implementation and the extensive tests mentioned in the proposal are contributed.