import os

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.placement_types import Partial, Replicate


class CustomMatmul(torch.autograd.Function):
    """Custom autograd.Function that does the same thing as torch.mm."""

    @staticmethod
    def forward(ctx, x, w):
        ctx.save_for_backward(x, w)
        return torch.mm(x, w)

    @staticmethod
    def backward(ctx, grad_output):
        x, w = ctx.saved_tensors
        return grad_output @ w.T, x.T @ grad_output


class TestModel(nn.Module):
    def __init__(self, use_transpose: bool, use_custom_fn: bool):
        super().__init__()
        # 3D weight to allow .transpose(1, 2)
        self.weight = nn.Parameter(torch.randn(4, 8, 8))
        self.use_transpose = use_transpose
        self.use_custom_fn = use_custom_fn
        self._grad_placement = None

    def forward(self, x):
        w = self.weight

        if isinstance(w, DTensor):
            grad_placements = [
                Partial() if isinstance(p, Replicate) else p for p in w.placements
            ]
            w = w.to_local(grad_placements=grad_placements)

            def _capture(param, model=self):
                if param.grad is not None and isinstance(param.grad, DTensor):
                    model._grad_placement = param.grad.placements

            self.weight.register_post_accumulate_grad_hook(_capture)

        if self.use_transpose:
            # .transpose() backward produces non-contiguous grad
            w = w.transpose(1, 2).contiguous()

        # Use first "expert" slice for matmul
        if self.use_custom_fn:
            return CustomMatmul.apply(x, w[0])
        else:
            return torch.mm(x, w[0])


def main():
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    torch.cuda.set_device(rank)

    # 1D mesh — parameters will be Replicate on this mesh
    mesh = init_device_mesh("cuda", (world_size,))

    torch.manual_seed(0)

    results = []
    for use_transpose in [False, True]:
        for use_custom_fn in [False, True]:
            t_label = "+transpose" if use_transpose else "           "
            fn_label = "custom_fn" if use_custom_fn else "native_mm"
            label = f"{fn_label} {t_label}"

            model = TestModel(
                use_transpose=use_transpose, use_custom_fn=use_custom_fn
            ).cuda()

            # Directly make weight a Replicate DTensor (no FSDP2)
            with torch.no_grad():
                model.weight = nn.Parameter(
                    DTensor.from_local(
                        model.weight.data,
                        device_mesh=mesh,
                        placements=[Replicate()],
                    )
                )

            x = torch.randn(2, 8, device="cuda", requires_grad=True)
            out = model(x)
            try:
                out.sum().backward()
                failed = False
            except AssertionError as e:
                failed = True
                if rank == 0:
                    print(f"  {label}: FAILED (assertion: {e})")

            p = model._grad_placement
            is_partial = (
                p is not None
                and all(isinstance(pl, Partial) for pl in p)
            )

            if rank == 0:
                if failed:
                    status = "FAILED (assertion)"
                elif is_partial:
                    status = f"ok   grad={p}"
                else:
                    status = f"BUG  grad={p}"
                print(f"  {label}: {status}")
                results.append((label, is_partial and not failed))

    if rank == 0:
        print()
        all_pass = all(ok for _, ok in results)
        if all_pass:
            print("All passed — bug not reproduced with upstream DTensor.")
        else:
            print("Bug reproduced with upstream DTensor:")
            print(".transpose() causes non-contiguous grad in backward,")
            print("which triggers Partial→Replicate redistribution in AccumulateGrad.")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

---

return std::move(new_grad
      .new_empty_strided_symint(
          variable.sym_sizes(),
          variable.sym_strides(),
          ...)
      .copy_(new_grad));