tensorflow_probability/python/internal/distribution_util.py
def move_dimension(x, source_idx, dest_idx):
  """Move a single tensor dimension within its shape.

  This is a special case of `tf.transpose()`, which applies
  arbitrary permutations to tensor dimensions.

  Args:
    x: Tensor of rank `ndims`.
    source_idx: Integer index into `x.shape` (negative indexing is supported).
    dest_idx: Integer index into `x.shape` (negative indexing is supported).

  Returns:
    x_perm: Tensor of rank `ndims`, in which the dimension at original
     index `source_idx` has been moved to new index `dest_idx`, with
     all other dimensions retained in their original order.

  Example:

  ```python
  x = tf.placeholder(shape=[200, 30, 4, 1, 6])
  x_perm = _move_dimension(x, 1, 1) # no-op
  x_perm = _move_dimension(x, 0, 3) # result shape [30, 4, 1, 200, 6]
  x_perm = _move_dimension(x, 0, -2) # equivalent to previous
  x_perm = _move_dimension(x, 4, 2) # result shape [200, 30, 6, 4, 1]
  ```
  """
  ndims = prefer_static_rank(x)
  dtype = dtype_util.common_dtype([source_idx, dest_idx],
                                  preferred_dtype=tf.int32)
  source_idx = tf.convert_to_tensor(value=source_idx, dtype=dtype)
  dest_idx = tf.convert_to_tensor(value=dest_idx, dtype=dtype)

  # Handle negative indexing.
  source_idx = pick_scalar_condition(source_idx < 0, ndims + source_idx,
                                     source_idx)
  dest_idx = pick_scalar_condition(dest_idx < 0, ndims + dest_idx, dest_idx)

  # Construct the appropriate permutation of dimensions, depending
  # whether the source is before or after the destination.
  def move_left_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, dest_idx, dtype=dtype), [source_idx],
            tf.range(dest_idx, source_idx, dtype=dtype),
            tf.range(source_idx + 1, ndims, dtype=dtype)
        ],
                  axis=0))

  def move_right_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, source_idx, dtype=dtype),
            tf.range(source_idx + 1, dest_idx + 1, dtype=dtype), [source_idx],
            tf.range(dest_idx + 1, ndims, dtype=dtype)
        ],
                  axis=0))

  def x_permuted():
    return tf.transpose(
        a=x,
        perm=prefer_static.cond(source_idx < dest_idx,
                                move_right_permutation,
                                move_left_permutation))

  # One final conditional to handle the special case where source
  # and destination indices are equal.
  return prefer_static.cond(tf.equal(source_idx, dest_idx),
                            lambda: x, x_permuted)
def move_dimension(x, source_idx, dest_idx):
  ndims = prefer_static_rank(x)
  dtype = dtype_util.common_dtype([source_idx, dest_idx],
                                  preferred_dtype=tf.int32)
  source_idx = tf.convert_to_tensor(value=source_idx, dtype=dtype)
  dest_idx = tf.convert_to_tensor(value=dest_idx, dtype=dtype)

  # Handle negative indexing.
  source_idx = pick_scalar_condition(source_idx < 0, ndims + source_idx,
                                     source_idx)
  dest_idx = pick_scalar_condition(dest_idx < 0, ndims + dest_idx, dest_idx)

  # Construct the appropriate permutation of dimensions, depending
  # whether the source is before or after the destination.
  def move_left_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, dest_idx, dtype=dtype), [source_idx],
            tf.range(dest_idx, source_idx, dtype=dtype),
            tf.range(source_idx + 1, ndims, dtype=dtype)
        ],
                  axis=0))

  def move_right_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, source_idx, dtype=dtype),
            tf.range(source_idx + 1, dest_idx + 1, dtype=dtype), [source_idx],
            tf.range(dest_idx + 1, ndims, dtype=dtype)
        ],
                  axis=0))

  def x_permuted():
    return tf.transpose(
        a=x,
        perm=prefer_static.cond(source_idx < dest_idx,
                                move_right_permutation,
                                move_left_permutation))

  # One final conditional to handle the special case where source
  # and destination indices are equal.
  return prefer_static.cond(tf.equal(source_idx, dest_idx),
                            lambda: x, x_permuted)
def move_dimension ( x , source_idx , dest_idx ) : ndims = prefer_static_rank ( x ) dtype = dtype_util . common_dtype ( [ source_idx , dest_idx ] , preferred_dtype = tf . int32 ) source_idx = tf . convert_to_tensor ( value = source_idx , dtype = dtype ) dest_idx = tf . convert_to_tensor ( value = dest_idx , dtype = dtype ) # Handle negative indexing. source_idx = pick_scalar_condition ( source_idx < 0 , ndims + source_idx , source_idx ) dest_idx = pick_scalar_condition ( dest_idx < 0 , ndims + dest_idx , dest_idx ) # Construct the appropriate permutation of dimensions, depending # whether the source is before or after the destination. def move_left_permutation ( ) : return prefer_static_value ( tf . concat ( [ tf . range ( 0 , dest_idx , dtype = dtype ) , [ source_idx ] , tf . range ( dest_idx , source_idx , dtype = dtype ) , tf . range ( source_idx + 1 , ndims , dtype = dtype ) ] , axis = 0 ) ) def move_right_permutation ( ) : return prefer_static_value ( tf . concat ( [ tf . range ( 0 , source_idx , dtype = dtype ) , tf . range ( source_idx + 1 , dest_idx + 1 , dtype = dtype ) , [ source_idx ] , tf . range ( dest_idx + 1 , ndims , dtype = dtype ) ] , axis = 0 ) ) def x_permuted ( ) : return tf . transpose ( a = x , perm = prefer_static . cond ( source_idx < dest_idx , move_right_permutation , move_left_permutation ) ) # One final conditional to handle the special case where source # and destination indices are equal. return prefer_static . cond ( tf . equal ( source_idx , dest_idx ) , lambda : x , x_permuted )
Move a single tensor dimension within its shape.

  This is a special case of `tf.transpose()`, which applies
  arbitrary permutations to tensor dimensions.

  Args:
    x: Tensor of rank `ndims`.
    source_idx: Integer index into `x.shape` (negative indexing is supported).
    dest_idx: Integer index into `x.shape` (negative indexing is supported).

  Returns:
    x_perm: Tensor of rank `ndims`, in which the dimension at original
     index `source_idx` has been moved to new index `dest_idx`, with
     all other dimensions retained in their original order.

  Example:

  ```python
  x = tf.placeholder(shape=[200, 30, 4, 1, 6])
  x_perm = _move_dimension(x, 1, 1) # no-op
  x_perm = _move_dimension(x, 0, 3) # result shape [30, 4, 1, 200, 6]
  x_perm = _move_dimension(x, 0, -2) # equivalent to previous
  x_perm = _move_dimension(x, 4, 2) # result shape [200, 30, 6, 4, 1]
  ```
Move a single tensor dimension within its shape .