code: def _make_expand_x_fn_for_batch_interpolation(y_ref, axis):
      """Make func to expand left/right (of axis) dims of tensors shaped like x."""
      # This expansion is to help x broadcast with `y`, the output.
      # In the batch case, the output shape is going to be
      #   Broadcast(y_ref.shape[:axis], x.shape[:-1]) +
      #   x.shape[-1:] +  y_ref.shape[axis+1:]
  
      # Recall we made axis non-negative
      y_ref_shape_left = ps.shape_slice(y_ref, np.s_[:axis])
      y_ref_shape_right = ps.shape_slice(y_ref, np.s_[axis + 1 :])
  
      def expand_right_dims(x, broadcast=False):
          """Expand x so it can bcast w/ tensors of output shape."""
          x_shape_left = ps.shape_slice(x, np.s_[:-1])
          x_shape_right = ps.shape_slice(x, np.s_[-1:])
          expanded_shape_left = ps.broadcast_shape(
              x_shape_left, _int32ones(ps.size(y_ref_shape_left))
          )
          expanded_shape = ps.concat(
              (
                  expanded_shape_left,
                  x_shape_right,
                  _int32ones(ps.size(y_ref_shape_right)),
              ),
              axis=0,
          )
          x_expanded = tf.reshape(x, expanded_shape)
          if broadcast:
              broadcast_shape_left = ps.broadcast_shape(
                  x_shape_left, y_ref_shape_left
              )
              ones_like_x_shape_right = _int32ones(1)
              broadcast_shape = ps.concat(
                  (
                      broadcast_shape_left,
                      ones_like_x_shape_right,
                      y_ref_shape_right,
                  ),
                  axis=0,
              )
              x_expanded = _broadcast_with(x_expanded, broadcast_shape)
          return x_expanded
  
      return expand_right_dims
  
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 11 for calling a global function with example usage: _make_expand_x_fn_for_batch_interpolation(y_ref, axis) and returns: expand_right_dims, x_expanded