import torch


def move_bdim_to_front(x, result_ndim=None):
    """
    Returns a tensor with a batch dimension at the front. If a batch
    dimension already exists, move it. Otherwise, create a new batch
    dimension at the front. If `result_ndim` is not None, ensure that the
    resulting tensor has rank equal to `result_ndim`.
    """
    x_dim = len(x.shape)
    x_bdim = x.bdim
    if x_bdim is None:
        x = torch.unsqueeze(x, 0)
    else:
        x = torch.movedim(x, x_bdim, 0)
    if result_ndim is None:
        return x
    diff = result_ndim - x_dim - (x_bdim is None)
    for _ in range(diff):
        x = torch.unsqueeze(x, 1)
    return x