tensorflow_probability/python/stats/quantiles.py
def _move_dims_to_flat_end(x, axis, x_ndims, right_end=True):
  """Move dims corresponding to `axis` in `x` to the end, then flatten.

  Args:
    x: `Tensor` with shape `[B0,B1,...,Bb]`.
    axis:  Python list of indices into dimensions of `x`.
    x_ndims:  Python integer holding number of dimensions in `x`.
    right_end:  Python bool.  Whether to move dims to the right end (else left).

  Returns:
    `Tensor` with value from `x` and dims in `axis` moved to end into one single
      dimension.
  """

  if not axis:
    return x

  # Suppose x.shape = [a, b, c, d]
  # Suppose axis = [1, 3]

  # other_dims = [0, 2] in example above.
  other_dims = sorted(set(range(x_ndims)).difference(axis))
  # x_permed.shape = [a, c, b, d]
  perm = other_dims + list(axis) if right_end else list(axis) + other_dims
  x_permed = tf.transpose(a=x, perm=perm)

  if x.shape.is_fully_defined():
    x_shape = x.shape.as_list()
    # other_shape = [a, c], end_shape = [b * d]
    other_shape = [x_shape[i] for i in other_dims]
    end_shape = [np.prod([x_shape[i] for i in axis])]
    full_shape = (
        other_shape + end_shape if right_end else end_shape + other_shape)
  else:
    other_shape = tf.gather(tf.shape(input=x), other_dims)
    full_shape = tf.concat(
        [other_shape, [-1]] if right_end else [[-1], other_shape], axis=0)
  return tf.reshape(x_permed, shape=full_shape)
def _move_dims_to_flat_end(x, axis, x_ndims, right_end=True):
  if not axis:
    return x

  # Suppose x.shape = [a, b, c, d]
  # Suppose axis = [1, 3]

  # other_dims = [0, 2] in example above.
  other_dims = sorted(set(range(x_ndims)).difference(axis))
  # x_permed.shape = [a, c, b, d]
  perm = other_dims + list(axis) if right_end else list(axis) + other_dims
  x_permed = tf.transpose(a=x, perm=perm)

  if x.shape.is_fully_defined():
    x_shape = x.shape.as_list()
    # other_shape = [a, c], end_shape = [b * d]
    other_shape = [x_shape[i] for i in other_dims]
    end_shape = [np.prod([x_shape[i] for i in axis])]
    full_shape = (
        other_shape + end_shape if right_end else end_shape + other_shape)
  else:
    other_shape = tf.gather(tf.shape(input=x), other_dims)
    full_shape = tf.concat(
        [other_shape, [-1]] if right_end else [[-1], other_shape], axis=0)
  return tf.reshape(x_permed, shape=full_shape)
def _move_dims_to_flat_end ( x , axis , x_ndims , right_end = True ) : if not axis : return x # Suppose x.shape = [a, b, c, d] # Suppose axis = [1, 3] # other_dims = [0, 2] in example above. other_dims = sorted ( set ( range ( x_ndims ) ) . difference ( axis ) ) # x_permed.shape = [a, c, b, d] perm = other_dims + list ( axis ) if right_end else list ( axis ) + other_dims x_permed = tf . transpose ( a = x , perm = perm ) if x . shape . is_fully_defined ( ) : x_shape = x . shape . as_list ( ) # other_shape = [a, c], end_shape = [b * d] other_shape = [ x_shape [ i ] for i in other_dims ] end_shape = [ np . prod ( [ x_shape [ i ] for i in axis ] ) ] full_shape = ( other_shape + end_shape if right_end else end_shape + other_shape ) else : other_shape = tf . gather ( tf . shape ( input = x ) , other_dims ) full_shape = tf . concat ( [ other_shape , [ - 1 ] ] if right_end else [ [ - 1 ] , other_shape ] , axis = 0 ) return tf . reshape ( x_permed , shape = full_shape )
Move dims corresponding to `axis` in `x` to the end, then flatten.

  Args:
    x: `Tensor` with shape `[B0,B1,...,Bb]`.
    axis:  Python list of indices into dimensions of `x`.
    x_ndims:  Python integer holding number of dimensions in `x`.
    right_end:  Python bool.  Whether to move dims to the right end (else left).

  Returns:
    `Tensor` with value from `x` and dims in `axis` moved to end into one single
      dimension.
Move dims corresponding to axis in x to the end then flatten .