Skip to content

Pull requests: jax-ml/jax

New pull request New
432 Open 17,371 Closed
432 Open 17,371 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
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[Mosaic GPU][NFC] Delete workaround for dialect bindings before jaxlib 0.5.1.
#26843 opened Feb 28, 2025 by copybara-service bot Loading…
@bchetioui
[mgpu] Forach in tiled layout.
#26842 opened Feb 28, 2025 by copybara-service bot Loading…
[Mosaic GPU] Fix as_dialect_barrier_memref to take into account BarrierRef's offset.
#26841 opened Feb 28, 2025 by copybara-service bot Loading…
@bchetioui
DOC: Remove duplicate JAX version 0.4.37 heading in Changelog
#26840 opened Feb 28, 2025 by rajasekharporeddy Loading…
Fixed printing order of results in jax.debug.print documentation
#26839 opened Feb 28, 2025 by Sai-Suraj-27 Loading…
1
Plumb layout through the creation of IFRT Arrays (roll-forward with fix).
#26838 opened Feb 28, 2025 by copybara-service bot Loading…
@emilyfertig
Fix thread safety of JAX error checking
#26837 opened Feb 28, 2025 by copybara-service bot Loading…
@ayaka14732
Update setup.py to automatically pick up libtpu patch releases
#26835 opened Feb 28, 2025 by skye Loading…
Add build targets for jax-rocm-plugin and jax-rocm-pjrt wheels.
#26827 opened Feb 27, 2025 by copybara-service bot Loading…
Print the list of installed packages before running pytests
#26826 opened Feb 27, 2025 by copybara-service bot Loading…
@nitins17
Added butterworth filter code
#26824 opened Feb 27, 2025 by ParamThakkar123 Loading…
1
Reimplement custom_vjp.optimize_remat using custom_dce pull ready Ready for copybara import and testing
#26814 opened Feb 27, 2025 by dfm Loading…
@dfm
Update ad.backward_pass to support non-linear functions of constants
#26811 opened Feb 27, 2025 by dfm Loading…
@dfm
6
Fix assertion fails on DebugInfo with JAX_USE_DIRECT_LINEARIZE=1
#26810 opened Feb 27, 2025 by gnecula Loading…
1
[Mosaic GPU] Add support for small RHS tile sizes in WGMMA
#26807 opened Feb 27, 2025 by copybara-service bot Loading…
@apaszke
Add a profiler test for gpu run
#26786 opened Feb 26, 2025 by copybara-service bot Loading…
Remove tensorstore dependency from test dependencies for python3.13-nogil jobs.
#26777 opened Feb 26, 2025 by copybara-service bot Loading…
[ROCm] Fix ROCm builds not finding numa library
#26775 opened Feb 26, 2025 by charleshofer Loading…
1
Add plugin versions to print_environment_info.
#26768 opened Feb 26, 2025 by dfm Loading…
@dfm
fix Initializer protocol
#26765 opened Feb 26, 2025 by Qwlouse Loading…
Add functionality to allow promoting RC wheels during release
#26756 opened Feb 26, 2025 by copybara-service bot Loading…
@nitins17
[Mosaic:TPU][Relayout] Row shifts for packed types and non-native tilings
#26754 opened Feb 25, 2025 by copybara-service bot Loading…
Debug failing CUDA test jobs CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#26749 opened Feb 25, 2025 by copybara-service bot Loading…
@nitins17
For testing jaxlib at head, install the libtpu version that jax pins
#26709 opened Feb 24, 2025 by copybara-service bot Loading…
@nitins17
Use the same name for aliased Vars when pretty-printing Jaxprs. pull ready Ready for copybara import and testing
#26697 opened Feb 24, 2025 by gnecula Loading…
@gnecula
ProTip! Type g i on any issue or pull request to go back to the issue listing page.