import torch


def permute_to_N_HWA_K(retina_tensor: torch.Tensor, num_classes: int) -> torch.Tensor:
    """ Transpose/reshape a tensor from (N, (A x K), H, W) to (N, (HxWxA), K). This is
    the tensor outputted from the RetinaNet head per pyramid level.
    Usage:
    >>> permute_to_N_HWA_K(torch.randn(1, 4 * 9, 4, 4), 4).shape
    torch.Size([1, 144, 4])
    """

    assert retina_tensor.dim() == 4, retina_tensor.shape
    N, _, H, W = retina_tensor.shape
    retina_tensor = retina_tensor.view(N, -1, num_classes, H, W)
    retina_tensor = retina_tensor.permute(0, 3, 4, 1, 2)
    retina_tensor = retina_tensor.reshape(N, -1, num_classes)  # Size=(N,HWA,K)
    return retina_tensor