Skip to content

Pull requests: jax-ml/jax

New pull request New
647 Open 20,400 Closed
647 Open 20,400 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

[ROCm] Fix GPU handle pool singleton aliasing using opaque handle typ…
#30715 opened Aug 3, 2025 by phambinhfin Loading…
[Mosaic TPU] Cleanup old checks
#30712 opened Aug 3, 2025 by copybara-service bot Loading…
Put NVIDIA wheels versions in build/nvidia-requirements.txt.
#30706 opened Aug 2, 2025 by copybara-service bot Loading…
[hijax] enable a Box test by adding a bit of qdd sensitivity pull ready Ready for copybara import and testing
#30705 opened Aug 2, 2025 by mattjj Loading…
@mattjj
[docs] New tutorial on Tracing
#30703 opened Aug 1, 2025 by melissawm Loading…
@emilyfertig
Fix minor code errors in Mosaic GPU docs.
#30702 opened Aug 1, 2025 by copybara-service bot Loading…
@jburnim
Upgrade CUDA version to 12.9.1, CUDNN version to 9.11.0.
#30701 opened Aug 1, 2025 by copybara-service bot Loading…
Add nvidia wheel versions in the lock files for python 3.14 and 3.14-ft.
#30700 opened Aug 1, 2025 by copybara-service bot Loading…
Run pytests on Windows with bigger JAX_NUM_GENERATED_CASES.
#30694 opened Aug 1, 2025 by copybara-service bot Loading…
[Mosaic GPU] Refine the layout used for SMEM<->GMEM transfers
#30683 opened Aug 1, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Use a single barrier for all loads in the WS pipeline
#30681 opened Aug 1, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Disable support for stationary inputs to WS pipelines and add a test
#30680 opened Aug 1, 2025 by copybara-service bot Loading…
@apaszke
Fix ravel_pytree to Return 1D Empty Array for Empty Pytrees (#30645)
#30664 opened Jul 31, 2025 by Aniketsy Loading…
2
Fix missing vmap_method in jax.debug.breakpoint
#30659 opened Jul 31, 2025 by patrick-kidger Loading…
Migrate dependencies on jax experimental targets to jax/experimental/BUILD.
#30652 opened Jul 31, 2025 by copybara-service bot Loading…
@danielsuo
Update stateful-computations.md link
#30646 opened Jul 31, 2025 by yufengg Loading…
[jax:windows] Fix windows x64 test failures by explicitly casting inputs.
#30640 opened Jul 31, 2025 by copybara-service bot Loading…
@danielsuo
#sdy fix JAX shmap returning pure_callback.
#30632 opened Jul 31, 2025 by copybara-service bot Loading…
1
[Draft][Mosaic:TPU][Relayouts] Vcombine retilings
#30628 opened Jul 31, 2025 by copybara-service bot Loading…
Test Bazel TPU postsubmit.
#30620 opened Jul 30, 2025 by copybara-service bot Loading…
An initial implementation of blocked map mma.
#30617 opened Jul 30, 2025 by copybara-service bot Loading…
[jaxlib] Fix namespace.
#30616 opened Jul 30, 2025 by copybara-service bot Loading…
@danielsuo
Include experimental/__init__.py in experimental targets so the jax.experimental submodules can be found.
#30615 opened Jul 30, 2025 by copybara-service bot Loading…
@danielsuo
[Pallas Fuser] Add Fuser rules for the repeat primitive.
#30612 opened Jul 30, 2025 by copybara-service bot Loading…
Modify pmap_test to use all devices. This is to enforce in underlying runtimes that all devices are used or a contiguous mesh with a valid subset of devices is used.
#30611 opened Jul 30, 2025 by copybara-service bot Loading…
ProTip! Adding no:label will show everything without a label.