jax

20efbd96 - [pallas:mosaic_gpu] Change the fori tests to also take the while_p path and fix the bug.

[pallas:mosaic_gpu] Change the fori tests to also take the while_p path and fix the bug. The bug was that bounds were dropped ctx.avals_in and then they were being extracted. Extract them before dropping them. PiperOrigin-RevId: 708266659