"""Minimal reproducer for fake PG all_gather_into_tensor failure.

The fake backend's _allgather_base uses output.chunk() internally,
creating views. When the input tensor has requires_grad=True, the
inplace write into those views triggers an autograd safety check.
Real NCCL doesn't hit this because it writes via a C++ kernel that
bypasses autograd.

Usage:
  # Fake PG (fails):
  python repro_fake_pg_allgather.py --backend fake
  # Real NCCL (works):
  torchrun --nproc_per_node 2 repro_fake_pg_allgather.py --backend nccl
"""
import argparse
import torch
import torch.distributed as dist


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--backend", choices=["fake", "nccl"], required=True)
    args = parser.parse_args()

    torch.cuda.set_device(0 if args.backend == "fake" else int(__import__("os").environ.get("LOCAL_RANK", 0)))

    if args.backend == "fake":
        dist.init_process_group("fake", rank=0, world_size=4)
    else:
        dist.init_process_group("nccl")

    world_size = dist.get_world_size()
    input_tensor = torch.randn(4, device="cuda", requires_grad=True)
    output_tensor = torch.empty(4 * world_size, device="cuda")

    try:
        dist.all_gather_into_tensor(output_tensor, input_tensor)
        print(f"backend={args.backend}, requires_grad=True: SUCCESS")
    except RuntimeError as e:
        print(f"backend={args.backend}, requires_grad=True: FAILED: {e}")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()