Immutable assignment (akin to array.at[...] in JAX) #159784
Description
🚀 The feature, motivation and pitch
torch.compile, torch.func and friends do not like mutable operations. However, sometimes assignment is the most concise way to implement an operation.
JAX has an immutable assignment operator (https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html). Could a similar feature be added to PyTorch tensors?
Alternatives
No response
Additional context
No response
cc @bdhirsh @ezyang @chauhang @penguinwu @zou3519 @Chillee @samdow @kshitij12345 @albanD