torch_harmonics.distributed.unpad_and_unflatten_leading_dims

torch_harmonics.distributed.unpad_and_unflatten_leading_dims(

tensor,

lead_shape,

lead_size,

num_trailing_dims=2,

)

Inverse of flatten_and_pad_leading_dims(): drop the padding rows and restore the leading dims.

The trailing num_trailing_dims dims are taken from tensor as-is, so this is valid even when the transform changed them (e.g. nlat, nlon -> lmax, mmax). num_trailing_dims must match the value passed to the flatten call.

Parameters:

• tensor (Tensor)

• lead_size (int)

• num_trailing_dims (int)

Return type:

Tensor