@tf.function
def grid_sample(moving, grid):
    'Given a moving image and a sampling grid as input, computes the\n    transformed image by sampling the moving image at locations given by\n    the grid.\n\n    Currently, only 2-D images, i.e., 4-D inputs are supported.\n\n    Parameters\n    ----------\n    moving : tf.Tensor, shape (N, H, W, C)\n        The moving image.\n    grid : tf.Tensor, shape (N, H, W, C)\n        A tensor of sampling points (x, y). The x and y values should be\n        normalized to [-1.0, 1.0] range.\n\n    Returns\n    -------\n    moved : tf.Tensor, shape (N, H, W, C)\n        The transformed image.\n\n    Notes\n    -----\n    Let M be the moving image of shape (H, W, C), T be the transformed\n    image of the same shape and G be the 2-D sampling grid of shape\n    (H, W, 2). The value of T at a location (x, y) is T[y, x, :] =\n    M[y\', x\', :] where [x\', y\'] = G[y, x, :].\n\n    Further, [x\', y\'] = [x + dx, y + dy] where [dx, dy] are the\n    displacements outputted by the CNN. When dx and dy are 0, the\n    sampling grid G is a regular grid and the transformed image is the\n    same as the moving image.\n\n    Since the sampling point (x + dx, y + dy) can be non-integral, the\n    value M[y\', x\'] is calculated using bi-linear interpolation.\n\n    References\n    ----------\n    .. [1] `Jaderberg, Max, Karen Simonyan, and Andrew Zisserman. "Spatial\n        transformer networks." Advances in neural information processing\n        systems. 2015. <https://arxiv.org/abs/1506.02025>`_\n    .. [2] `TensorFlow implementation of spatial transformer networks.\n        <https://github.com/tensorflow/models/tree/master/research/transformer>`_\n    .. [3] `Spatial Transformer Networks by Kushagra Bhatnagar\n        <https://link.medium.com/0b2OrmqVO5>`_\n    '
    (nb, nh, nw, nc) = moving.shape
    x = grid[(..., 0)]
    y = grid[(..., 1)]
    x = tf.cast(x, 'float32')
    y = tf.cast(y, 'float32')
    x = (((x + 1.0) * 0.5) * tf.cast(nw, 'float32'))
    y = (((y + 1.0) * 0.5) * tf.cast(nh, 'float32'))
    y_max = tf.cast((nh - 1), 'int32')
    x_max = tf.cast((nw - 1), 'int32')
    zero = tf.constant(0, 'int32')
    x0 = tf.cast(tf.floor(x), 'int32')
    x1 = (x0 + 1)
    y0 = tf.cast(tf.floor(y), 'int32')
    y1 = (y0 + 1)
    x0 = tf.clip_by_value(x0, zero, x_max)
    x1 = tf.clip_by_value(x1, zero, x_max)
    y0 = tf.clip_by_value(y0, zero, y_max)
    y1 = tf.clip_by_value(y1, zero, y_max)
    b = (tf.ones_like(x0) * tf.reshape(tf.range(nb), [nb, 1, 1]))
    idx_a = tf.stack([b, y0, x0], axis=(- 1))
    idx_b = tf.stack([b, y1, x0], axis=(- 1))
    idx_c = tf.stack([b, y0, x1], axis=(- 1))
    idx_d = tf.stack([b, y1, x1], axis=(- 1))
    moving_a = tf.gather_nd(moving, idx_a)
    moving_b = tf.gather_nd(moving, idx_b)
    moving_c = tf.gather_nd(moving, idx_c)
    moving_d = tf.gather_nd(moving, idx_d)
    x0_f = tf.cast(x0, 'float32')
    x1_f = tf.cast(x1, 'float32')
    y0_f = tf.cast(y0, 'float32')
    y1_f = tf.cast(y1, 'float32')
    wa = tf.expand_dims(((x1_f - x) * (y1_f - y)), axis=(- 1))
    wb = tf.expand_dims(((x1_f - x) * (y - y0_f)), axis=(- 1))
    wc = tf.expand_dims(((x - x0_f) * (y1_f - y)), axis=(- 1))
    wd = tf.expand_dims(((x - x0_f) * (y - y0_f)), axis=(- 1))
    moved = tf.add_n([(wa * moving_a), (wb * moving_b), (wc * moving_c), (wd * moving_d)])
    return moved
Given a moving image and a sampling grid as input, computes the
transformed image by sampling the moving image at locations given by
the grid.

Currently, only 2-D images, i.e., 4-D inputs are supported.

Parameters
----------
moving : tf.Tensor, shape (N, H, W, C)
    The moving image.
grid : tf.Tensor, shape (N, H, W, C)
    A tensor of sampling points (x, y). The x and y values should be
    normalized to [-1.0, 1.0] range.

Returns
-------
moved : tf.Tensor, shape (N, H, W, C)
    The transformed image.

Notes
-----
Let M be the moving image of shape (H, W, C), T be the transformed
image of the same shape and G be the 2-D sampling grid of shape
(H, W, 2). The value of T at a location (x, y) is T[y, x, :] =
M[y', x', :] where [x', y'] = G[y, x, :].

Further, [x', y'] = [x + dx, y + dy] where [dx, dy] are the
displacements outputted by the CNN. When dx and dy are 0, the
sampling grid G is a regular grid and the transformed image is the
same as the moving image.

Since the sampling point (x + dx, y + dy) can be non-integral, the
value M[y', x'] is calculated using bi-linear interpolation.

References
----------
.. [1] `Jaderberg, Max, Karen Simonyan, and Andrew Zisserman. "Spatial
    transformer networks." Advances in neural information processing
    systems. 2015. <https://arxiv.org/abs/1506.02025>`_
.. [2] `TensorFlow implementation of spatial transformer networks.
    <https://github.com/tensorflow/models/tree/master/research/transformer>`_
.. [3] `Spatial Transformer Networks by Kushagra Bhatnagar
    <https://link.medium.com/0b2OrmqVO5>`_
@tf.function
def grid_sample(moving, grid):
    'Given a moving image and a sampling grid as input, computes the\n    transformed image by sampling the moving image at locations given by\n    the grid.\n\n    Currently, only 2-D images, i.e., 4-D inputs are supported.\n\n    Parameters\n    ----------\n    moving : tf.Tensor, shape (N, H, W, C)\n        The moving image.\n    grid : tf.Tensor, shape (N, H, W, C)\n        A tensor of sampling points (x, y). The x and y values should be\n        normalized to [-1.0, 1.0] range.\n\n    Returns\n    -------\n    moved : tf.Tensor, shape (N, H, W, C)\n        The transformed image.\n\n    Notes\n    -----\n    Let M be the moving image of shape (H, W, C), T be the transformed\n    image of the same shape and G be the 2-D sampling grid of shape\n    (H, W, 2). The value of T at a location (x, y) is T[y, x, :] =\n    M[y\', x\', :] where [x\', y\'] = G[y, x, :].\n\n    Further, [x\', y\'] = [x + dx, y + dy] where [dx, dy] are the\n    displacements outputted by the CNN. When dx and dy are 0, the\n    sampling grid G is a regular grid and the transformed image is the\n    same as the moving image.\n\n    Since the sampling point (x + dx, y + dy) can be non-integral, the\n    value M[y\', x\'] is calculated using bi-linear interpolation.\n\n    References\n    ----------\n    .. [1] `Jaderberg, Max, Karen Simonyan, and Andrew Zisserman. "Spatial\n        transformer networks." Advances in neural information processing\n        systems. 2015. <https://arxiv.org/abs/1506.02025>`_\n    .. [2] `TensorFlow implementation of spatial transformer networks.\n        <https://github.com/tensorflow/models/tree/master/research/transformer>`_\n    .. [3] `Spatial Transformer Networks by Kushagra Bhatnagar\n        <https://link.medium.com/0b2OrmqVO5>`_\n    '
    (nb, nh, nw, nc) = moving.shape
    x = grid[(..., 0)]
    y = grid[(..., 1)]
    x = tf.cast(x, 'float32')
    y = tf.cast(y, 'float32')
    x = (((x + 1.0) * 0.5) * tf.cast(nw, 'float32'))
    y = (((y + 1.0) * 0.5) * tf.cast(nh, 'float32'))
    y_max = tf.cast((nh - 1), 'int32')
    x_max = tf.cast((nw - 1), 'int32')
    zero = tf.constant(0, 'int32')
    x0 = tf.cast(tf.floor(x), 'int32')
    x1 = (x0 + 1)
    y0 = tf.cast(tf.floor(y), 'int32')
    y1 = (y0 + 1)
    x0 = tf.clip_by_value(x0, zero, x_max)
    x1 = tf.clip_by_value(x1, zero, x_max)
    y0 = tf.clip_by_value(y0, zero, y_max)
    y1 = tf.clip_by_value(y1, zero, y_max)
    b = (tf.ones_like(x0) * tf.reshape(tf.range(nb), [nb, 1, 1]))
    idx_a = tf.stack([b, y0, x0], axis=(- 1))
    idx_b = tf.stack([b, y1, x0], axis=(- 1))
    idx_c = tf.stack([b, y0, x1], axis=(- 1))
    idx_d = tf.stack([b, y1, x1], axis=(- 1))
    moving_a = tf.gather_nd(moving, idx_a)
    moving_b = tf.gather_nd(moving, idx_b)
    moving_c = tf.gather_nd(moving, idx_c)
    moving_d = tf.gather_nd(moving, idx_d)
    x0_f = tf.cast(x0, 'float32')
    x1_f = tf.cast(x1, 'float32')
    y0_f = tf.cast(y0, 'float32')
    y1_f = tf.cast(y1, 'float32')
    wa = tf.expand_dims(((x1_f - x) * (y1_f - y)), axis=(- 1))
    wb = tf.expand_dims(((x1_f - x) * (y - y0_f)), axis=(- 1))
    wc = tf.expand_dims(((x - x0_f) * (y1_f - y)), axis=(- 1))
    wd = tf.expand_dims(((x - x0_f) * (y - y0_f)), axis=(- 1))
    moved = tf.add_n([(wa * moving_a), (wb * moving_b), (wc * moving_c), (wd * moving_d)])
    return moved