Skip to content

Pull requests: jax-ml/jax

New pull request New
679 Open 23,403 Closed
679 Open 23,403 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Add test for buffer_protocol with donation
#34180 opened Jan 6, 2026 by copybara-service bot Loading…
@pschuh
Fix(cpu_client): preserve exception context in pure_callback cache hits
#34179 opened Jan 6, 2026 by nandan2003 Loading…
3 tasks done
1
3
Finalize deprecation of jax.lax.zeros_like_array
#34177 opened Jan 6, 2026 by copybara-service bot Loading…
@jakevdp
Automated Code Change
#34176 opened Jan 6, 2026 by copybara-service bot Loading…
Raise a better error message when a Tracer is closed over
#34174 opened Jan 6, 2026 by copybara-service bot Loading…
@yashk2810
[ci] temporarily pin numpy<2.4 for array API tests pull ready Ready for copybara import and testing
#34173 opened Jan 6, 2026 by jakevdp Loading…
@jakevdp
1
[pallas:sc] Allowed specifying tiling in pltpu.emit_pipeline
#34171 opened Jan 6, 2026 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Artificially limit SMEM capacity for tests
#34169 opened Jan 6, 2026 by copybara-service bot Loading…
@apaszke
Automated Code Change
#34167 opened Jan 6, 2026 by copybara-service bot Loading…
Fixes Typo in Pseudocode Example
#34164 opened Jan 6, 2026 by JyotinderSingh Loading…
1
[graphcast] Prepare for jax_pmap_shmap_merge=True.
#34163 opened Jan 6, 2026 by copybara-service bot Loading…
@danielsuo
[hijax] start implementing custom_vjp on top of hijax primitives kokoro:force-run pull ready Ready for copybara import and testing
#34161 opened Jan 6, 2026 by mattjj Loading…
3
Switch CCCL headers version to 2.8.3 downloaded from Github.
#34160 opened Jan 6, 2026 by copybara-service bot Loading…
@ybaturina
1
Plumb prim params for call discharge rule (to handle named_computation_p
#34154 opened Jan 5, 2026 by copybara-service bot Loading…
@sharadmv
Plumb prim params for call discharge rule (to handle named_computation_p
#34150 opened Jan 5, 2026 by copybara-service bot Loading…
@sharadmv
Add on-demand TPU profiling tests to profiler_test.py.
#34146 opened Jan 5, 2026 by copybara-service bot Loading…
[pallas:mosaic] Lower lax.cond to scf.index_switch
#34144 opened Jan 5, 2026 by copybara-service bot Loading…
@superbobry
clarify usage and guarantees of jax.clear_caches() documentation
#34141 opened Jan 4, 2026 by adityasingh345 Loading…
1
@mattjj
4
improve aot stages.Compiled.call pytree mismatch error message better_errors Improve the error reporting
#34140 opened Jan 3, 2026 by mattjj • Draft
@mattjj
4
Docs: clarify NumPy PRNG comparison in random numbers guide
#34138 opened Jan 3, 2026 by amarjeetsingh146 Loading…
1
@jakevdp
12
Fixes minor typos in Pallas documentation
#34137 opened Jan 3, 2026 by JyotinderSingh Loading…
1
[ROCm] update to test if there are GPU events when doing profiling on…
#34135 opened Jan 2, 2026 by cj401-amd Loading…
4
Automated Code Change
#34131 opened Jan 2, 2026 by copybara-service bot Loading…
Improve __repr__ for custom_jvp-wrapped functions
#34130 opened Jan 2, 2026 by Utkarshkarki Loading…
14
Shorten splash attention kernel name to avoid 255 char limit
#34125 opened Dec 30, 2025 by patrick-toulme Loading…
4
ProTip! Follow long discussions with comments:>50.