"""Minimal repro: DTensor batch norm backward hangs with channel-dim sharding.

Without a registered strategy for aten.native_batch_norm_backward.default,
the decomposition sharding propagator enters an infinite loop in to_meta()
when inputs are Shard(1).

Usage: torchrun --nproc_per_node=2 repro_batchnorm_hang.py
"""
import torch
import torch.nn.functional as F
from torch.distributed.tensor import Shard, distribute_tensor, init_device_mesh
import torch.distributed as dist

dist.init_process_group("nccl")
torch.cuda.set_device(dist.get_rank())
mesh = init_device_mesh("cuda", (dist.get_world_size(),))

C = dist.get_world_size() * 2
x = distribute_tensor(
    torch.randn(2, C, 3, 3, device="cuda").requires_grad_(True), mesh, [Shard(1)]
)
w = distribute_tensor(torch.randn(C, device="cuda").requires_grad_(True), mesh, [Shard(0)])
b = distribute_tensor(torch.randn(C, device="cuda").requires_grad_(True), mesh, [Shard(0)])
rm = distribute_tensor(torch.zeros(C, device="cuda"), mesh, [Shard(0)])
rv = distribute_tensor(torch.ones(C, device="cuda"), mesh, [Shard(0)])

out = F.batch_norm(x, rm, rv, w, b, training=True)  # forward works
out.sum().backward()  # hangs here
print(f"[rank {dist.get_rank()}] backward done", flush=True)

dist.destroy_process_group()