torch_harmonics.distributed.polar_halo_exchange

torch_harmonics.distributed.polar_halo_exchange(x, r_lat)

Exchange r_lat halo rows with neighbouring polar ranks.

Gathers r_lat latitude rows from each polar neighbour and returns a halo-padded tensor. Boundary ranks receive zero-padding on the missing side. The operation is fully differentiable: the backward pass sends halo gradients back to their owning ranks and accumulates them.

Parameters:

• x (torch.Tensor) – Input tensor of shape (B, C, H_local, W).

• r_lat (int) – Number of halo rows to exchange on each side.

Returns:

Halo-padded tensor of shape (B, C, H_local + 2 * r_lat, W).

Return type:

torch.Tensor