jax.lax.transpose

jax.lax.transpose(operand, permutation)[source]

Wraps XLA’s Transpose operator.