jax

b47fa65c - Merge pull request #30040 from mattjj:direct-linearize-fixes-5