Pull requests: jax-ml/jax
Pull requests list
Dedup prev_outs and cur_inps in looped_opt_barrier. As you can see below, we get rid of the redundant opt_barriers.
#40337
by copybara-service
Bot
was merged Aug 31, 2026
Loading…
If the cur_eqn is program_order_p (or any HOP in the future) add an opt_barrier for every
cur_inp i.e. cur_inps = [opt_barrier((token, c))[1] for c in cur_inps] where token, prev_outs_new = optimization_barrier((token, prev_outs)). This is strictly better than the approach before this change (described below).
#40336
by copybara-service
Bot
was merged Aug 31, 2026
Loading…
Fix tsan/asan build breakage after https://github.com/jax-ml/jax/commit/75b9bd22d4f5b76165a7a733c18664ae8800ad37
#40334
by copybara-service
Bot
was merged Aug 31, 2026
Loading…
[CI] Skip PallasMGPUTest.test_store_reshaped_ref.
#40332
by copybara-service
Bot
was merged Aug 31, 2026
Loading…
Use portserver for locally-run multiprocess CPU tests.
#40330
by copybara-service
Bot
was merged Aug 31, 2026
Loading…
Add jax::ft_mutex and jax::ft_lock_guard with thread-safety annotations.
#40328
by copybara-service
Bot
was merged Aug 31, 2026
Loading…
Initialize arg_exclude_mask with the correct values if exclude_argnames is None
#40323
by copybara-service
Bot
was merged Aug 30, 2026
Loading…
Update broken profiling docs links
kokoro:force-run
pull ready
Ready for copybara import and testing
#40322
by austin362667
was merged Aug 31, 2026
Loading…
[pallas] Optional hashable tag in
pltpu.get_barrier_semaphore(tag=...).
#40316
by copybara-service
Bot
was merged Sep 1, 2026
Loading…
Use program_order_p even if exclude_argnames is None
#40313
by copybara-service
Bot
was merged Aug 29, 2026
Loading…
Explicitly propagate and specify strip_memory_space in state discharge callers.
#40311
by copybara-service
Bot
was merged Aug 28, 2026
Loading…
Support mutable operations (
set, swap, addupdate) on ReshapeTransform and BitcastTransform views during state discharge.
#40310
by copybara-service
Bot
was merged Aug 28, 2026
Loading…
[remat3] re-implement handling of checkpoint_name in custom_vjp fwd
kokoro:force-run
pull ready
Ready for copybara import and testing
#40309
by mattjj
Collaborator
was closed Aug 29, 2026
Loading…
[pallas:fuser] Support full-dimension (dynamic_)slice and (dynamic_)update_slice
#40308
by copybara-service
Bot
was merged Aug 28, 2026
Loading…
[docs] more tweaks to complex number description
documentation
kokoro:force-run
pull ready
Ready for copybara import and testing
#40306
by mattjj
Collaborator
was merged Aug 28, 2026
Loading…
Use partition_list and merge_list in program_order
#40302
by copybara-service
Bot
was closed Aug 29, 2026
Loading…
Add
exclude_argnames on program_order.
#40298
by copybara-service
Bot
was merged Aug 28, 2026
Loading…
Exclude Windows Python 3.15t from the pytest build matrix.
#40296
by copybara-service
Bot
was merged Aug 28, 2026
Loading…
ProTip!
Type g i on any issue or pull request to go back to the issue listing page.