import torch


def sum_over_all_but_batch_and_last_n(
    tensor: torch.Tensor, n_dims: int
) -> torch.Tensor:
    """
    Returns the sum of the input tensor over all dimensions except
    the first (batch) and last n_dims.

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

    Returns:
        New tensor of shape (B, X[0], X[1], ..., X[n_dims-1]).
        Will return the unchanged input tensor if `tensor.dim() == n_dims + 1`

    Examples:
        import torch

        A = torch.ones(2,3,4)
        print(sum_over_all_but_batch_and_last_n(A, 1))
        # prints torch.Size([2, 4])
        print(sum_over_all_but_batch_and_last_n(A, 2))
        # prints torch.Size([2, 3, 4])
    """
    if tensor.dim() == n_dims + 1:
        return tensor
    else:
        dims = list(range(1, tensor.dim() - n_dims))
        return tensor.sum(dim=dims)