jax

6259cfe3 - Fold transposes feeding the RHS of a matmul into a change in the dot dimension numbers.