jax

e96b28c4 - [Pallas] Add missing trace_stop ops for jaxprs that end without a non-scoped op.