tensorflow_probability/python/math/interpolation.py
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 = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[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 = tf.pad(
        tensor=tf.shape(input=x),
        paddings=[[axis, tf.size(input=y_ref_shape_right)]],
        constant_values=1)
    x_expanded = tf.reshape(x, expanded_shape)

    if broadcast:
      out_shape = tf.concat((
          y_ref_shape_left,
          tf.shape(input=x),
          y_ref_shape_right,
      ),
                            axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.cast(tf.zeros(out_shape), tf.bool)
      else:
        x_expanded += tf.zeros(out_shape, dtype=x.dtype)
    return x_expanded

  return expand_ends
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 = tf.shape(input=y_ref)
  y_ref_shape_left = y_ref_shape[:axis]
  y_ref_shape_right = y_ref_shape[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 = tf.pad(
        tensor=tf.shape(input=x),
        paddings=[[axis, tf.size(input=y_ref_shape_right)]],
        constant_values=1)
    x_expanded = tf.reshape(x, expanded_shape)

    if broadcast:
      out_shape = tf.concat((
          y_ref_shape_left,
          tf.shape(input=x),
          y_ref_shape_right,
      ),
                            axis=0)
      if x.dtype.is_bool:
        x_expanded = x_expanded | tf.cast(tf.zeros(out_shape), tf.bool)
      else:
        x_expanded += tf.zeros(out_shape, dtype=x.dtype)
    return x_expanded

  return expand_ends
Make func to expand left/right (of axis) dims of tensors shaped like x.