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
def _make_expand_x_fn_for_non_batch_interpolation ( y_ref , axis ) : # 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.
Make func to expand left / right ( of axis ) dims of tensors shaped like x .