willychan21/ParallelKernelBench_Problems

1from typing import Optional, Sequence2 3import torch4import torch.distributed as dist5 6 7def _all_to_all_transpose(8    tensor: torch.Tensor,9    split_dim: int,10    group: dist.ProcessGroup,11) -> list[torch.Tensor]:12    world_size = dist.get_world_size(group)13    chunk = tensor.shape[split_dim] // world_size14    send = [x.contiguous() for x in torch.split(tensor, chunk, dim=split_dim)]15    recv = [torch.empty_like(send[0]) for _ in range(world_size)]16    dist.all_to_all(recv, send, group=group)17    return recv18 19 20def _truncate(tensor: torch.Tensor, dim: int, size: int) -> torch.Tensor:21    slices = [slice(None)] * tensor.ndim22    slices[dim % tensor.ndim] = slice(0, size)23    return tensor[tuple(slices)].contiguous()24 25 26@torch.no_grad()27def solution(28    x: torch.Tensor,29    s: Sequence[int],30    dim: Sequence[int],31    norm: str = "ortho",32    group: Optional[dist.ProcessGroup] = None,33) -> torch.Tensor:34    group = group or dist.group.WORLD35    dim0, dim1 = int(dim[0]), int(dim[1])36 37    n0 = int(s[0]) if s[0] is not None else None38    n1 = int(s[1]) if s[1] is not None else None39 40    x1 = torch.fft.fft(x, n=n0, dim=dim0, norm=norm)41 42    x1_recv = _all_to_all_transpose(x1, split_dim=dim0, group=group)43    x1_tran = torch.cat(x1_recv, dim=dim1)44 45    x2 = torch.fft.fft(x1_tran, n=n1, dim=dim1, norm=norm)46 47    return _truncate(x2, dim1, x2.shape[dim1] // 2 + 1)