pytorch - 💡(How to fix) Fix `torch.compile` backward crashes when compiled function output is an expanded tensor

RuntimeError: unsupported operation: more than one element of the written-to
tensor refers to a single memory location. Please clone() the tensor before
performing the operation.

---

import torch

def fn(x):
    return x.sum(dim=0, keepdim=True).expand_as(x)

x = torch.randn(4, 8, device="cuda", requires_grad=True)

# Eager: works
fn(x).sum().backward()
print(f"Eager grad: {x.grad.shape}")  # torch.Size([4, 8])

# Compiled: crashes
x2 = x.detach().clone().requires_grad_(True)
torch.compile(fn, backend="inductor")(x2).sum().backward()

---

RuntimeError: unsupported operation: more than one element of the written-to
tensor refers to a single memory location. Please clone() the tensor before
performing the operation.

---

# ALL reductions + expand → CRASH
x.sum(dim=0, keepdim=True).expand_as(x)
x.mean(dim=0, keepdim=True).expand_as(x)
x.amax(dim=0, keepdim=True).expand_as(x)
x.amin(dim=0, keepdim=True).expand_as(x)
x.prod(dim=0, keepdim=True).expand_as(x)
x.norm(dim=0, keepdim=True).expand_as(x)
x.std(dim=0, keepdim=True).expand_as(x)
x.var(dim=0, keepdim=True).expand_as(x)
x.logsumexp(dim=0, keepdim=True).expand_as(x)

# Non-reduction computation + expand → ALSO CRASH
x[:1].sin().expand_as(x)
x[:1].softmax(-1).expand_as(x)
(x[:1] * 2 + 1).expand_as(x)

# broadcast_to (equivalent to expand) → CRASH
x.sum(0, keepdim=True).broadcast_to(4, 8)

# Any dim, 2D/3D/ND:
x.sum(dim=1, keepdim=True).expand_as(x)        # 2D, CRASH
x.sum(dim=0, keepdim=True).expand_as(x)        # 3D shape (4,8,16), CRASH

---

# Pure view of input + expand: OK (AOT Autograd correctly handles input aliasing)
x[:1].expand_as(x)                     # OK
x.narrow(0, 0, 1).expand_as(x)         # OK

# Expand followed by any materializing op: OK
x.sum(0, keepdim=True).expand_as(x).sin()         # OK
x.sum(0, keepdim=True).expand_as(x) * x           # OK
x.sum(0, keepdim=True).expand_as(x) + 0           # OK
x.sum(0, keepdim=True).expand_as(x).contiguous()  # OK
x.sum(0, keepdim=True).expand_as(x).clone()       # OK

# repeat (copies memory, no aliasing): OK
x.sum(0, keepdim=True).repeat(4, 1)    # OK

# Trivial expand (size 1 → size 1, no actual aliasing): OK
x = torch.randn(1, 8, requires_grad=True)
x.sum(0, keepdim=True).expand_as(x)    # OK (expand is no-op)