jax

eb95074f - Merge pull request #36346 from mattjj:remove-split-transpose