code: def _batch_interp_with_gather_nd(
      x, x_idx_unclipped, y_ref, nd, fill_value, batch_ndims
  ):
      """Batch interpolation starting with indices."""
      dtype = x.dtype
      # Wherever x is NaN, x_idx_unclipped will be NaN as well.
      # Keep track of the nan indices here (so we can impute NaN later).
      # Also eliminate any NaN indices, since there is not NaN in 32bit.
      nan_idx = tf.math.is_nan(x_idx_unclipped)
      x_idx_unclipped = tf.where(
          nan_idx, tf.cast(0.0, dtype=dtype), x_idx_unclipped
      )
  
      # ny[k] is number of y reference points in interp dim k.
      # It is used to indicate the dimension sizes.
      ny = tf.cast(
          ps.shape_slice(y_ref, np.s_[batch_ndims : batch_ndims + nd]),
          dtype,
      )
  
      # x_idx.shape = [A1, ..., An, D, nd]
      x_idx = tf.clip_by_value(
          x_idx_unclipped, tf.zeros((), dtype=dtype), ny - 1
      )
  
      # Get the index above and below x_idx.
      # Naively we could set idx_below = floor(x_idx), idx_above = ceil(x_idx),
      # however, this results in idx_below == idx_above whenever x is on a grid.
      # This in turn results in y_ref_below == y_ref_above, and then the gradient
      # at this point is zero.  So here we 'jitter' one of idx_below, idx_above,
      # so that they are at different values.  This jittering does not affect the
      # interpolated value, but does make the gradient nonzero (unless of course
      # the y_ref values are the same).
      idx_below = tf.floor(x_idx)
      idx_above = tf.minimum(idx_below + 1, ny - 1)
      idx_below = tf.maximum(idx_above - 1, 0)
  
      # These are the values of y_ref corresponding to above/below indices.
      # idx_below_int32.shape = x.shape[:-1] + [nd]
      idx_below_int32 = tf.cast(idx_below, dtype=tf.int32)
      idx_above_int32 = tf.cast(idx_above, dtype=tf.int32)
  
      # idx_below_list is a length nd list of shape x.shape[:-1] int32 tensors.
      idx_below_list = tf.unstack(idx_below_int32, axis=-1)
      idx_above_list = tf.unstack(idx_above_int32, axis=-1)
  
      # Use t to get a convex combination of the below/above values.
      # t.shape = [A1, ..., An, D, nd]
      t = x_idx - idx_below
  
      # x, and tensors shaped like x, need to be added to, and selected with
      # (using tf.where) the output y.  This requires appending singletons.
      def _expand_x_fn(tensor):
          # Reshape tensor to tensor.shape + [1] * M.
          extended_shape = ps.concat(
              [
                  ps.shape(tensor),
                  ps.ones_like(
                      ps.convert_to_shape_tensor(
                          ps.shape_slice(
                              y_ref, np.s_[batch_ndims + nd :]
                          )
                      )
                  ),
              ],
              axis=0,
          )
          return tf.reshape(tensor, extended_shape)
  
      # Now, t.shape = [A1, ..., An, D, nd] + [1] * (rank(y_ref) - nd - batch_ndims)
      t = _expand_x_fn(t)
      s = 1 - t
  
      # Re-insert NaN wherever x was NaN.
      nan_idx = _expand_x_fn(nan_idx)
      t = tf.where(nan_idx, tf.constant(np.nan, dtype), t)
  
      # Initialize y and accumulate in a loop. An alternative would be to store
      # summands in a list. However, without XLA compilation, the "list method"
      # results in storage of 2^nd (possibly large) summands, which could OOM.
      # Thus, if you are not XLA compiling, the method below is highly preferred.
      # With XLA compilation, both methods are equivalent.
      y = tf.zeros((), dtype=dtype)
  
      # Our work above has located x's fractional index inside a cube of above/below
      # indices. The distance to the below indices is t, and to the above indices
      # is s.
      # Drawing lines from x to the cube walls, we get 2**nd smaller cubes. Each
      # term in the result is a product of a reference point, gathered from y_ref,
      # multiplied by a volume.  The volume is that of the cube opposite to the
      # reference point.  E.g. if the reference point is below x in every axis, the
      # volume is that of the cube with corner above x in every axis, s[0]*...*s[nd]
      # We could probably do this with one massive gather, but that would be very
      # unreadable and un-debuggable.  It also would create a large Tensor.
      for zero_ones_list in _binary_count(nd):
          gather_from_y_ref_idx = []
          opposite_volume_t_idx = []
          opposite_volume_s_idx = []
          for k, zero_or_one in enumerate(zero_ones_list):
              if zero_or_one == 0:
                  # If the kth iterate has zero_or_one = 0,
                  # Will gather from the 'below' reference point along axis k.
                  gather_from_y_ref_idx.append(idx_below_list[k])
                  # Now append the index to gather for computing opposite_volume.
                  # This could be done by initializing opposite_volume to 1, then here:
                  #  opposite_volume *= tf.gather(s, indices=k, axis=tf.rank(x) - 1)
                  # but that puts a gather in the 'inner loop.'  Better to append the
                  # index and do one larger gather down below.
                  opposite_volume_s_idx.append(k)
              else:
                  gather_from_y_ref_idx.append(idx_above_list[k])
                  # Append an index to gather, having the same effect as
                  #   opposite_volume *= tf.gather(t, indices=k, axis=tf.rank(x) - 1)
                  opposite_volume_t_idx.append(k)
  
          # Compute opposite_volume (volume of cube opposite the ref point):
          # Recall t.shape = s.shape = [D, nd] + [1, ..., 1]
          # Gather from t and s along the 'nd' axis, which is rank(x) - 1.
          ov_axis = ps.cast(ps.rank(x) - 1, tf.int32)
          opposite_volume = tf.reduce_prod(
              tf.gather(
                  t,
                  indices=tf.cast(
                      opposite_volume_t_idx, dtype=tf.int32
                  ),
                  axis=ov_axis,
              ),
              axis=ov_axis,
          ) * tf.reduce_prod(
              tf.gather(
                  s,
                  indices=tf.cast(
                      opposite_volume_s_idx, dtype=tf.int32
                  ),
                  axis=ov_axis,
              ),
              axis=ov_axis,
          )  # pyformat: disable
  
          y_ref_pt = tf.gather_nd(
              y_ref,
              tf.stack(gather_from_y_ref_idx, axis=-1),
              batch_dims=batch_ndims,
          )
  
          y = y + y_ref_pt * opposite_volume
  
      if tf.debugging.is_numeric_tensor(fill_value):
          # Recall x_idx_unclipped.shape = [D, nd],
          # so here we check if it was out of bounds in any of the nd dims.
          # Thus, oob_idx.shape = [D].
          oob_idx = tf.reduce_any(
              (x_idx_unclipped < 0) | (x_idx_unclipped > ny - 1),
              axis=-1,
          )
  
          # Now, y.shape = [D, B1,...,BM], so we'll have to broadcast oob_idx.
  
          oob_idx = _expand_x_fn(oob_idx)  # Shape [D, 1,...,1]
          oob_idx = _broadcast_with(oob_idx, ps.shape(y))
          y = tf.where(oob_idx, fill_value, y)
      return y
  
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 6 for calling a global function with example usage: _batch_interp_with_gather_nd(x, x_idx_unclipped, y_ref, nd, fill_value, batch_ndims) and returns: y, tf