I wrote the following function. Your comments would be much appreciated:
def interp1d(y: torch.Tensor, newx: torch.Tensor): 
    '''
    Function for simple linear interpolation in pytorch. 
    Assumes x is [0, 1, 2, 3, ..., len(y)-1]
    newx can be any dimension and the dimension of result will match newx's. 

    Test: 
    y = torch.arange(24, dtype=torch.float)
    newx = torch.tensor([1.3, 2.5])
    newx = torch.arange(24).reshape(2,3,4) + 0.23
    newx = torch.arange(24).reshape(2,3,4) 
    '''
    assert len(y.shape) == 1 # 1-dimensional
    assert len(y) > 1
    assert newx.min() >= 0
    assert newx.max() <= len(y) - 1

    npoints = len(y)
    ndim = len(newx.shape)

    x = newx.unsqueeze(ndim)
    x0 = torch.arange(npoints-1)
    diff = (x - x0).to(torch.float)
    bin = (diff >= 0) & (diff < 1)

    y0 = y.to(torch.float)
    y0left = y[:-1].expand(diff.shape)
    y0right = y[1:].expand(diff.shape)
    y0max = y0[-1].expand(diff.shape) / (npoints - 1) # Averaging because it gets summed (npoints - 1) times later.

    weight = diff * bin
    interp = torch.lerp(y0left, y0right, weight) * bin
    interp = interp + (x == (npoints - 1)) * y0max # The corner case
    
    result = interp.sum(dim=ndim)

    return result