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,
)[source]#

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.

Complex inputs (the forward SHTs call this on their spectral output) are processed as their real view: this is pure layout code touching only the leading dims, and a copy over a complex buffer cannot be codegen’d by inductor – triton has no complex type, so it fails with KeyError: 'complex64'. view_as_real appends the re/im pair as a trailing dim, which is simply carried along; the contiguous() right before the view_as_complex supplies the unit last-dim stride the latter requires.

Parameters:
Return type:

Tensor