def _batch_transpose(mat):
n = distribution_util.prefer_static_rank(mat)<EOL>perm = tf.range(n)<EOL>perm = tf.concat([perm[:-<NUM_LIT:2>], [perm[-<NUM_LIT:1>], perm[-<NUM_LIT:2>]]], axis=<NUM_LIT:0>)<EOL>return tf.transpose(a=mat, perm=perm)<EOL>
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.