code: def batch_interp_regular_nd_grid(
      x,
      x_ref_min,
      x_ref_max,
      y_ref,
      axis,
      fill_value="constant_extension",
      name=None,
  ):
      """Multi-linear interpolation on a regular (constant spacing) grid.
  
      Given [a batch of] reference values, this function computes a multi-linear
      interpolant and evaluates it on [a batch of] of new `x` values. This is a
      multi-dimensional generalization of [Bilinear Interpolation](
      https://en.wikipedia.org/wiki/Bilinear_interpolation).
  
      The interpolant is built from reference values indexed by `nd` dimensions
      of `y_ref`, starting at `axis`.
  
      The x grid span is defined by `x_ref_min`, `x_ref_max`. The number of grid
      points is inferred from the shape of `y_ref`.
  
      For example, take the case of a `2-D` scalar valued function and no leading
      batch dimensions.  In this case, `y_ref.shape = [C1, C2]` and `y_ref[i, j]`
      is the reference value corresponding to grid point
  
      ```
      [x_ref_min[0] + i * (x_ref_max[0] - x_ref_min[0]) / (C1 - 1),
       x_ref_min[1] + j * (x_ref_max[1] - x_ref_min[1]) / (C2 - 1)]
      ```
  
      In the general case, dimensions to the left of `axis` in `y_ref` are broadcast
      with leading dimensions in `x`, `x_ref_min`, `x_ref_max`.
  
      Args:
        x: Numeric `Tensor` The x-coordinates of the interpolated output values for
          each batch.  Shape `[..., D, nd]`, designating [a batch of] `D`
          coordinates in `nd` space.  `D` must be `>= 1` and is not a batch dim.
        x_ref_min:  `Tensor` of same `dtype` as `x`.  The minimum values of the
          (implicitly defined) reference `x_ref`.  Shape `[..., nd]`.
        x_ref_max:  `Tensor` of same `dtype` as `x`.  The maximum values of the
          (implicitly defined) reference `x_ref`.  Shape `[..., nd]`.
        y_ref:  `Tensor` of same `dtype` as `x`.  The reference output values. Shape
          `[..., C1, ..., Cnd, B1,...,BM]`, designating [a batch of] reference
          values indexed by `nd` dimensions, of a shape `[B1,...,BM]` valued
          function (for `M >= 0`).
        axis:  Scalar integer `Tensor`.  Dimensions `[axis, axis + nd)` of `y_ref`
          index the interpolation table.  E.g. `3-D` interpolation of a scalar
          valued function requires `axis=-3` and a `3-D` matrix valued function
          requires `axis=-5`.
        fill_value:  Determines what values output should take for `x` values that
          are below `x_ref_min` or above `x_ref_max`. Scalar `Tensor` or
          'constant_extension' ==> Extend as constant function.
          Default value: `'constant_extension'`
        name:  A name to prepend to created ops.
          Default value: `'batch_interp_regular_nd_grid'`.
  
      Returns:
        y_interp:  Interpolation between members of `y_ref`, at points `x`.
          `Tensor` of same `dtype` as `x`, and shape `[..., D, B1, ..., BM].`
  
      Exceptions will be raised if shapes are statically determined to be wrong.
  
      Raises:
        ValueError:  If `rank(x) < 2`.
        ValueError:  If `axis` is not a scalar.
        ValueError:  If `axis + nd > rank(y_ref)`.
  
      #### Examples
  
      Interpolate a function of one variable.
  
      ```python
      y_ref = tf.exp(tf.linspace(start=0., stop=10., num=20))
  
      tfp.math.batch_interp_regular_nd_grid(
          # x.shape = [3, 1], x_ref_min/max.shape = [1].  Trailing `1` for `1-D`.
          x=[[6.0], [0.5], [3.3]], x_ref_min=[0.], x_ref_max=[10.], y_ref=y_ref,
          axis=0)
      ==> approx [exp(6.0), exp(0.5), exp(3.3)]
      ```
  
      Interpolate a scalar function of two variables.
  
      ```python
      x_ref_min = [0., 0.]
      x_ref_max = [2 * np.pi, 2 * np.pi]
  
      # Build y_ref.
      x0s, x1s = tf.meshgrid(
          tf.linspace(x_ref_min[0], x_ref_max[0], num=100),
          tf.linspace(x_ref_min[1], x_ref_max[1], num=100),
          indexing='ij')
  
      def func(x0, x1):
        return tf.sin(x0) * tf.cos(x1)
  
      y_ref = func(x0s, x1s)
  
      x = 2 * np.pi * tf.random.uniform(shape=(10, 2))
  
      tfp.math.batch_interp_regular_nd_grid(x, x_ref_min, x_ref_max, y_ref, axis=-2)
      ==> tf.sin(x[:, 0]) * tf.cos(x[:, 1])
      ```
  
      """
      with tf.name_scope(name or "batch_interp_regular_nd_grid"):
          dtype = dtype_util.common_dtype(
              [x, x_ref_min, x_ref_max, y_ref], dtype_hint=tf.float32
          )
  
          # Arg checking.
          fill_value = _intake_fill_value_for_nd_interp(
              fill_value, dtype
          )
  
          # x.shape = [..., nd].
          x = tf.convert_to_tensor(x, name="x", dtype=dtype)
          _assert_ndims_statically(x, expect_ndims_at_least=2)
  
          # y_ref.shape = [..., C1,...,Cnd, B1,...,BM]
          y_ref = tf.convert_to_tensor(y_ref, name="y_ref", dtype=dtype)
  
          # x_ref_min.shape = [nd]
          x_ref_min = tf.convert_to_tensor(
              x_ref_min, name="x_ref_min", dtype=dtype
          )
          x_ref_max = tf.convert_to_tensor(
              x_ref_max, name="x_ref_max", dtype=dtype
          )
          _assert_ndims_statically(
              x_ref_min, expect_ndims_at_least=1, expect_static=True
          )
          _assert_ndims_statically(
              x_ref_max, expect_ndims_at_least=1, expect_static=True
          )
  
          # nd is the number of dimensions indexing the interpolation table, it's the
          # 'nd' in the function name.
          nd = tf.compat.dimension_value(x_ref_min.shape[-1])
          if nd is None:
              raise ValueError(
                  "`x_ref_min.shape[-1]` must be known statically."
              )
          tensorshape_util.assert_is_compatible_with(
              x_ref_max.shape[-1:], x_ref_min.shape[-1:]
          )
  
          # Convert axis and check it statically.
          axis = _intake_axis_for_nd_interp(axis, y_ref, nd)
  
          x_batch_shape = ps.shape_slice(x, np.s_[:-2])
          x_ref_min_batch_shape = ps.shape_slice(x_ref_min, np.s_[:-1])
          x_ref_max_batch_shape = ps.shape_slice(x_ref_max, np.s_[:-1])
          y_ref_batch_shape = ps.shape_slice(y_ref, np.s_[:axis])
  
          # Do a brute-force broadcast of batch dims (add zeros).
          batch_shape = y_ref_batch_shape
          for tensor in [
              x_batch_shape,
              x_ref_min_batch_shape,
              x_ref_max_batch_shape,
          ]:
              batch_shape = ps.broadcast_shape(batch_shape, tensor)
  
          def _batch_shape_of_zeros_with_rightmost_singletons(
              n_singletons,
          ):
              """Return Tensor of zeros with some singletons on the rightmost dims."""
              return ps.concat(
                  [batch_shape, _int32ones(n_singletons)], axis=0
              )
  
          x = _broadcast_with(
              x,
              _batch_shape_of_zeros_with_rightmost_singletons(
                  n_singletons=2
              ),
          )
          x_ref_min = _broadcast_with(
              x_ref_min,
              _batch_shape_of_zeros_with_rightmost_singletons(
                  n_singletons=1
              ),
          )
          x_ref_max = _broadcast_with(
              x_ref_max,
              _batch_shape_of_zeros_with_rightmost_singletons(
                  n_singletons=1
              ),
          )
          y_ref = _broadcast_with(
              y_ref,
              _batch_shape_of_zeros_with_rightmost_singletons(
                  n_singletons=ps.rank(y_ref) - axis
              ),
          )
  
          # At this point,
          # x.shape = [A1, ..., An, D, nd], where n = batch_ndims
          # and
          # y_ref.shape = [A1, ..., An, C1, C2,..., Cnd, B1,...,BM]
          # y_ref[A1, ..., An, i1,...,ind] is a shape [B1,...,BM] Tensor with value
          # at index [i1,...,ind] in the interpolation table.
          #  and x_ref_max have shapes [A1, ..., An, nd].
  
          batch_ndims = ps.rank(x) - 2
  
          # ny[k] is number of y reference points in interp dim k.
          # It is used to indicate the dimension sizes.
          ny = tf.cast(
              # After broadcasting y_ref with x, slice(batch_ndims, batch_ndims + nd)
              # is the proper way to extract ny. Before broadcasting, use
              # slice(axis, axis + nd)
              ps.shape_slice(
                  y_ref, np.s_[batch_ndims : batch_ndims + nd]
              ),
              dtype,
          )
  
          # Map [x_ref_min, x_ref_max] to [0, ny - 1].
          # This is the (fractional) index of x.
          # x_idx_unclipped[A1, ..., An, d, k] is the fractional index into dim k of
          # interpolation table for the dth x value.
          x_ref_min_expanded = tf.expand_dims(x_ref_min, axis=-2)
          x_ref_max_expanded = tf.expand_dims(x_ref_max, axis=-2)
          x_idx_unclipped = (
              (ny - 1)
              * (x - x_ref_min_expanded)
              / (x_ref_max_expanded - x_ref_min_expanded)
          )
  
          return _batch_interp_with_gather_nd(
              x=x,
              x_idx_unclipped=x_idx_unclipped,
              y_ref=y_ref,
              nd=nd,
              fill_value=fill_value,
              batch_ndims=batch_ndims,
          )
  
imports: import itertools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 4 for calling a global function with example usage: batch_interp_regular_nd_grid(x, x_ref_min, x_ref_max, y_ref, axis, fill_value, name) and returns: _batch_interp_with_gather_nd, ps