import torch


def flatten(x: torch.Tensor, keepdims: bool = False):
    """
    Flattens B C H W input to B C*H*W output, optionally retains trailing dimensions.
    """
    y = x.view(x.size(0), -1)
    if keepdims:
        for _ in range(y.dim(), x.dim()):
            y = y.unsqueeze(-1)
    return y