def _make_expand_x_fn_for_batch_interpolation(y_ref, axis):
  """"""
  # 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 = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[axis + 1:]

  def expand_right_dims(x, broadcast=False):
    """Expand x so it can bcast w/ tensors of output shape."""
    expanded_shape_left = tf.broadcast_dynamic_shape(
        tf.shape(input=x)[:-1],
        tf.ones([tf.size(input=y_ref_shape_left)], dtype=tf.int32))
    expanded_shape = tf.concat(
        (expanded_shape_left, tf.shape(input=x)[-1:],
         tf.ones([tf.size(input=y_ref_shape_right)], dtype=tf.int32)),
        axis=0)
    x_expanded = tf.reshape(x, expanded_shape)
    if broadcast:
      broadcast_shape_left = tf.broadcast_dynamic_shape(
          tf.shape(input=x)[:-1], y_ref_shape_left)
      broadcast_shape = tf.concat(
          (broadcast_shape_left, tf.shape(input=x)[-1:], y_ref_shape_right),
          axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.cast(tf.zeros(broadcast_shape), tf.bool)
      else:
        x_expanded += tf.zeros(broadcast_shape, dtype=x.dtype)
    return x_expanded

  return expand_right_dims