def _pad_dims(tensor, rank):
    """Append empty dimensions to the tensor until it is of the given rank."""
    for _ in range(rank - tensor.shape.ndims):
        tensor = tensor[..., None]
    return tensor
