jax

b00a3a19 - Merge pull request #27015 from mattjj:direct-linearize-fixes-4