code: def _make_expand_x_fn_for_non_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 non-batch case, the output shape is going to be
      #   y_ref.shape[:axis] + x.shape + 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_ends(x, broadcast=False):
          """Expand x so it can bcast w/ tensors of output shape."""
          # Assume out_shape = A + x.shape + B, and rank(A) = axis.
          # Expand with singletons with same rank as A, B.
          expanded_shape = ps.concat(
              [ps.shape(x), _int32ones(ps.size(y_ref_shape_right))],
              axis=0,
          )
          x_expanded = tf.reshape(x, expanded_shape)
  
          if broadcast:
              out_bcaster = ps.concat(
                  (
                      y_ref_shape_left,
                      _int32ones(ps.rank(x)),
                      y_ref_shape_right,
                  ),
                  axis=0,
              )
              x_expanded = _broadcast_with(x_expanded, out_bcaster)
          return x_expanded
  
      return expand_ends
  
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 10 for calling a global function with example usage: _make_expand_x_fn_for_non_batch_interpolation(y_ref, axis) and returns: expand_ends, x_expanded