import torch


def get_spatial_meshgrid(x: torch.Tensor, scale=False):
    """
    Get grid which contains spatial coordinates at each pixel location

    Args:
        x: image of shape [batch_size, channels, height, width] for which
            we want to generate the spatial grid
    """
    batch_size, _, height, width = x.size()
    # Generate mesh grid
    xx = torch.arange(0, width).view(1, -1).repeat(height, 1)
    yy = torch.arange(0, height).view(-1, 1).repeat(1, width)
    xx = xx.view(1, 1, height, width).repeat(batch_size, 1, 1, 1)
    yy = yy.view(1, 1, height, width).repeat(batch_size, 1, 1, 1)
    grid = torch.cat((xx, yy), 1).float()
    if x.is_cuda:
        grid = grid.cuda()
    if scale:
        grid[:, 0] = grid[:, 0] / width
        grid[:, 1] = grid[:, 1] / height
    return grid