import torch


def _interp_array(start: torch.Tensor, end: torch.Tensor,
                  num_steps: int) -> torch.Tensor:
    """Linearly interpolate 2D tensors, returns 3D tensors.
    Args:
      start: 2D tensor for start point of interpolation of shape [x,y].
      end: 2D tensor as end point of interpolation of shape [x,y] (same as start).
      num_steps: number of steps to interpolate.
    Returns:
      New tensor of shape [num_steps, x, y]
    """
    alpha = torch.linspace(0., 1., num_steps)
    beta = 1 - alpha
    return torch.einsum('a,bc->abc', alpha, end) + torch.einsum('a,bc->abc', beta,
                                                          start)