torch_harmonics.distributed.polar_halo_reduce#

torch_harmonics.distributed.polar_halo_reduce(y, r_lat, nlat_local)[source]#

Accumulate r_lat halo rows onto their owning polar ranks.

The adjoint of polar_halo_exchange(), and its mirror image: where the exchange borrows rows from the neighbours before a gather-shaped operator reads them, this returns rows to the neighbours after a scatter-shaped operator has written them. Boundary ranks drop the outward halo, matching the zero-padding the exchange applies there. Fully differentiable; the backward is the exchange.

Parameters:

y (torch.Tensor) – Halo-padded tensor of shape (B, C, nlat_local + 2 * r_lat, W).
r_lat (int) – Number of halo rows on each side.
nlat_local (int) – Number of latitudes this rank owns, i.e. the extent of the returned tensor.

Returns:

Tensor of shape (B, C, nlat_local, W) holding this rank’s completed rows.

Return type:

torch.Tensor