@tf.function
def regular_grid(shape):
    'Returns a batch of 2-D regular grids.\n\n    Currently, only 2-D regular grids are supported.\n\n    Parameters\n    ----------\n    shape : sequence of ints, shape (3, )\n        The desired regular grid shape of the form (N, H, W).\n\n    Returns\n    -------\n    grid : tf.Tensor, shape (N, H, W, 2)\n        A batch of 2-D regular grids, values normalized to [-1.0, 1.0]\n        range.\n\n    Notes\n    -----\n    Sampling using the regular grid is an identity transformation, i.e.,\n    it results in the same input and output images.\n\n    References\n    ----------\n    .. [1] `NumPy, "numpy.meshgrid"\n        <https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html>`_\n    .. [2] `NumPy, "numpy.indices"\n        <https://numpy.org/doc/stable/reference/generated/numpy.indices.html>`_\n    '
    (nb, nh, nw) = shape
    x = tf.linspace((- 1.0), 1.0, nw)
    y = tf.linspace((- 1.0), 1.0, nh)
    (X, Y) = tf.meshgrid(x, y)
    grid = tf.stack([X, Y], axis=(- 1))
    grid = tf.expand_dims(grid, axis=0)
    multiples = tf.constant([nb, 1, 1, 1], tf.int32)
    grid = tf.tile(grid, multiples)
    return grid
Returns a batch of 2-D regular grids.

Currently, only 2-D regular grids are supported.

Parameters
----------
shape : sequence of ints, shape (3, )
    The desired regular grid shape of the form (N, H, W).

Returns
-------
grid : tf.Tensor, shape (N, H, W, 2)
    A batch of 2-D regular grids, values normalized to [-1.0, 1.0]
    range.

Notes
-----
Sampling using the regular grid is an identity transformation, i.e.,
it results in the same input and output images.

References
----------
.. [1] `NumPy, "numpy.meshgrid"
    <https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html>`_
.. [2] `NumPy, "numpy.indices"
    <https://numpy.org/doc/stable/reference/generated/numpy.indices.html>`_
@tf.function
def regular_grid(shape):
    'Returns a batch of 2-D regular grids.\n\n    Currently, only 2-D regular grids are supported.\n\n    Parameters\n    ----------\n    shape : sequence of ints, shape (3, )\n        The desired regular grid shape of the form (N, H, W).\n\n    Returns\n    -------\n    grid : tf.Tensor, shape (N, H, W, 2)\n        A batch of 2-D regular grids, values normalized to [-1.0, 1.0]\n        range.\n\n    Notes\n    -----\n    Sampling using the regular grid is an identity transformation, i.e.,\n    it results in the same input and output images.\n\n    References\n    ----------\n    .. [1] `NumPy, "numpy.meshgrid"\n        <https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html>`_\n    .. [2] `NumPy, "numpy.indices"\n        <https://numpy.org/doc/stable/reference/generated/numpy.indices.html>`_\n    '
    (nb, nh, nw) = shape
    x = tf.linspace((- 1.0), 1.0, nw)
    y = tf.linspace((- 1.0), 1.0, nh)
    (X, Y) = tf.meshgrid(x, y)
    grid = tf.stack([X, Y], axis=(- 1))
    grid = tf.expand_dims(grid, axis=0)
    multiples = tf.constant([nb, 1, 1, 1], tf.int32)
    grid = tf.tile(grid, multiples)
    return grid