import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import distribute_tensor, Shard

def main():
    dist.init_process_group(backend='gloo')
    rank = dist.get_rank()
    mesh = init_device_mesh('cpu', (2,))

    x = torch.arange(4).reshape(1, 4).float()
    dt = distribute_tensor(x, mesh, [Shard(0)])

    print(f'[rank {rank}] BEFORE: DTensor shape={dt.shape}, local={dt._local_tensor.shape}')
    dt.squeeze_(0)
    print(f'[rank {rank}] AFTER:  DTensor shape={dt.shape}, local={dt._local_tensor.shape}')

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

if op_call == aten.squeeze_.dim:
    args[0]._spec = output_spec
    return return_and_correct_aliasing(op_call, args, kwargs, args[0])

dt._local_tensor.squeeze_(0)  # Works - shape becomes [4]