jax

155e11f0 - Added changes from PR #626 and PR #645. This also fixes merge conflicts from upstream.

Added changes from PR #626 and PR #645. This also fixes merge conflicts from upstream. This uses the following commits: - 61625381cc2f014972e9a67692f3249e6ca473d9 - 9344b2b2a0827236b1d11feb63a7c6dbc4286124