Skip to content

Pull requests: jax-ml/jax

New pull request New
435 Open 17,365 Closed
435 Open 17,365 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

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 linux python 3.13t nightly tests
#26833 opened Feb 28, 2025 by copybara-service bot Loading…
Set -n to 32 to see if it fixes Pytest CPU failures
#26832 opened Feb 27, 2025 by copybara-service bot Loading…
@nitins17
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
[Mosaic GPU] Remove TMA inputs
#26806 opened Feb 27, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU][NFC] Start refactoring the MMA parameter inference
#26805 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…
[Mosaic GPU][NFC] Move some functions to a new file called inference_utils.py.
#26776 opened Feb 26, 2025 by copybara-service bot Loading…
@bchetioui
[ROCm] Fix ROCm builds not finding numa library
#26775 opened Feb 26, 2025 by charleshofer Loading…
1
[Pallas/Mosaic GPU] Add an abstraction to obtain a slice of dynamic shared memory when using waprgroup semantics.
#26773 opened Feb 26, 2025 by copybara-service bot Loading…
@bchetioui
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…
ProTip! Mix and match filters to narrow down what you’re looking for.