import torch
from typing import Iterable


def match_to(
    x: torch.Tensor, ref: torch.Tensor, keep_axes: Iterable[int] = (1,)
) -> torch.Tensor:
    """Modify the shape of an input to match a reference by repeating elements.

    Uses .expand() on the input.

    Args:
        x: Input tensor.
        ref: Reference tensor.
        keep_axes: Don't change these axes.

    Returns:
        The modified input.

    """

    target_shape = list(ref.shape)
    for i in keep_axes:
        target_shape[i] = x.shape[i]
    target_shape = tuple(target_shape)
    if x.shape == target_shape:
        pass
    elif x.ndim == 1:
        x = x.unsqueeze(0)
    else:
        while x.ndim < len(target_shape):
            x = x.unsqueeze(-1)

    x = x.expand(*target_shape)
    x = x.to(device=ref.device, dtype=ref.dtype)

    return x