code: 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)
  
file_path: probability/tensorflow_probability/python/optimizer/bfgs.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.optimizer import bfgs_utils
label_desc: Use image node_id 7 for calling a global function with example usage: _batch_transpose(mat) and returns: tf