tensorflow_probability/python/math/interpolation.py
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 = 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
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 = 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
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
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 .