jax

25d62ff5 - Add support for transposing tensor being loaded

Add support for transposing tensor being loaded This helps when the result is being used transposed. Without it when we go to LDS, we incur significant number of bank conflicts.