def discrete_sample(x, grid):
    """
    Args:
        x (Tensor): [N, C, H, W]
        grid (Tensor): [N, grid_H, grid_W, 2]
    Returns:
        output (Tensor): [N, C, grid_H, grid_W]
    """
    N, C, H, W = x.shape
    _, grid_H, grid_W, _ = grid.shape
    spatial_shape = paddle.to_tensor([[W, H]], dtype=paddle.float32)
    index = (grid * spatial_shape + 0.5).astype(paddle.int64).flatten(1, 2)
    h_index = index[:, :, 1].clip(0, H - 1)
    w_index = index[:, :, 0].clip(0, W - 1)
    batch_index = paddle.arange(N).unsqueeze(-1).tile([1, grid_H * grid_W])
    output = x[batch_index, :, h_index, w_index]
    output = output.transpose([0, 2, 1]).reshape([N, C, grid_H, grid_W])
    return output

    Args:
        x (Tensor): [N, C, H, W]
        grid (Tensor): [N, grid_H, grid_W, 2]
    Returns:
        output (Tensor): [N, C, grid_H, grid_W]
    
ppdet/modeling/transformers/utils.py