keras/src/backend/tensorflow/numpy.py/swapaxes
def swapaxes(x, axis1, axis2):
    x = convert_to_tensor(x)

    if (
        x.shape.rank is not None
        and isinstance(axis1, int)
        and isinstance(axis2, int)
    ):
        # This branch makes sure `perm` is statically known, to avoid a
        # not-compile-time-constant XLA error.
        axis1 = canonicalize_axis(axis1, x.ndim)
        axis2 = canonicalize_axis(axis2, x.ndim)

        # Directly return x if no movement is required
        if axis1 == axis2:
            return x

        perm = list(range(x.ndim))
        perm[axis1] = axis2
        perm[axis2] = axis1
    else:
        x_rank = tf.rank(x)
        axis1 = tf.where(axis1 < 0, tf.add(axis1, x_rank), axis1)
        axis2 = tf.where(axis2 < 0, tf.add(axis2, x_rank), axis2)
        perm = tf.range(x_rank)
        perm = tf.tensor_scatter_nd_update(
            perm, [[axis1], [axis2]], [axis2, axis1]
        )
    return tf.transpose(x, perm)