import torch


def resampler_wrapper_torch(imgs, coords):
    """
    equivalent to tfa.image.resampler
    Args:
    imgs: [N, H, W, C] images to resample
    coords: [N, H, W, 2], source pixel locations from which to copy
    Returns:
    [N, H, W, C] sampled pixels
    """
    return torch.nn.functional.grid_sample(
        imgs.permute([0, 3, 1, 2]),  # change images from (N, H, W, C) to (N, C, H, W)
        torch.tensor([-1, -1], device=imgs.device) + 2.0 * coords,  # resize coords from (0,1) to (-1, 1)
        align_corners=True
    ).permute([0, 2, 3, 1])