tensorflow_probability/python/optimizer/bfgs.py
def _batch_transpose(mat):
  """Transpose a possibly batched matrix.

  Args:
    mat: A `tf.Tensor` of shape `[..., n, m]`.

  Returns:
    A tensor of shape `[..., m, n]` with matching batch dimensions.
  """
  n = distribution_util.prefer_static_rank(mat)
  perm = tf.range(n)
  perm = tf.concat([perm[:-2], [perm[-1], perm[-2]]], axis=0)
  return tf.transpose(a=mat, perm=perm)
def _batch_transpose(mat):
  n = distribution_util.prefer_static_rank(mat)
  perm = tf.range(n)
  perm = tf.concat([perm[:-2], [perm[-1], perm[-2]]], axis=0)
  return tf.transpose(a=mat, perm=perm)
def _batch_transpose ( mat ) : n = distribution_util . prefer_static_rank ( mat ) perm = tf . range ( n ) perm = tf . concat ( [ perm [ : - 2 ] , [ perm [ - 1 ] , perm [ - 2 ] ] ] , axis = 0 ) return tf . transpose ( a = mat , perm = perm )
Transpose a possibly batched matrix.

  Args:
    mat: A `tf.Tensor` of shape `[..., n, m]`.

  Returns:
    A tensor of shape `[..., m, n]` with matching batch dimensions.
Transpose a possibly batched matrix .