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):
  """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)
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.