jax

de633956 - [jax] Add a test that runs reduction on host

[jax] Add a test that runs reduction on host Check that nested computations generated by offloaded computation are correctly outlined into the host module. PiperOrigin-RevId: 647771541