def sum_over_all_but_batch_and_last_n(tensor: torch.Tensor, n_dims: int) -> torch.Tensor:
    '\n    Calculates the sum over all dimensions, except the first\n    (batch dimension), and excluding the last n_dims.\n\n    This function will ignore the first dimension and it will\n    not aggregate over the last n_dims dimensions.\n\n    Args:\n        tensor: An input tensor of shape ``(B, ..., X[n_dims-1])``.\n        n_dims: Number of dimensions to keep.\n\n    Example:\n        >>> tensor = torch.ones(1, 2, 3, 4, 5)\n        >>> sum_over_all_but_batch_and_last_n(tensor, n_dims=2).shape\n        torch.Size([1, 4, 5])\n\n    Returns:\n        A tensor of shape ``(B, ..., X[n_dims-1])``\n    '
    if (tensor.dim() == (n_dims + 1)):
        return tensor
    else:
        dims = list(range(1, (tensor.dim() - n_dims)))
        return tensor.sum(dim=dims)
Calculates the sum over all dimensions, except the first
(batch dimension), and excluding the last n_dims.

This function will ignore the first dimension and it will
not aggregate over the last n_dims dimensions.

Args:
    tensor: An input tensor of shape ``(B, ..., X[n_dims-1])``.
    n_dims: Number of dimensions to keep.

Example:
    >>> tensor = torch.ones(1, 2, 3, 4, 5)
    >>> sum_over_all_but_batch_and_last_n(tensor, n_dims=2).shape
    torch.Size([1, 4, 5])

Returns:
    A tensor of shape ``(B, ..., X[n_dims-1])``
def sum_over_all_but_batch_and_last_n(tensor: torch.Tensor, n_dims: int) -> torch.Tensor:
    '\n    Calculates the sum over all dimensions, except the first\n    (batch dimension), and excluding the last n_dims.\n\n    This function will ignore the first dimension and it will\n    not aggregate over the last n_dims dimensions.\n\n    Args:\n        tensor: An input tensor of shape ``(B, ..., X[n_dims-1])``.\n        n_dims: Number of dimensions to keep.\n\n    Example:\n        >>> tensor = torch.ones(1, 2, 3, 4, 5)\n        >>> sum_over_all_but_batch_and_last_n(tensor, n_dims=2).shape\n        torch.Size([1, 4, 5])\n\n    Returns:\n        A tensor of shape ``(B, ..., X[n_dims-1])``\n    '
    if (tensor.dim() == (n_dims + 1)):
        return tensor
    else:
        dims = list(range(1, (tensor.dim() - n_dims)))
        return tensor.sum(dim=dims)