Skip to content

Pull requests: jax-ml/jax

New pull request New
635 Open 23,211 Closed
635 Open 23,211 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 a JAX config to disable the preemption service.
#33903 opened Dec 11, 2025 by copybara-service bot Loading…
@mwhittaker
Change VLOG to LOG(ERROR) in CUDA diagnostics.
#33901 opened Dec 11, 2025 by copybara-service bot Loading…
Enable async oneDNN build by default
#33900 opened Dec 11, 2025 by akhilgoe Loading…
1
Add device dict support to Pallas TPU interpret mode
#33898 opened Dec 11, 2025 by copybara-service bot Loading…
@sharadmv
improve nondiff argnums error when a tracer is passed pull ready Ready for copybara import and testing
#33897 opened Dec 11, 2025 by mattjj • Draft
@mattjj
3
Reverts 64a8e0de42681bdc26b89dba30fc2c979869ea82
#33896 opened Dec 11, 2025 by copybara-service bot Loading…
@ybaturina
1
[MGPU] Add WGStridedFragLayout broadcast support on major dim.
#33893 opened Dec 11, 2025 by copybara-service bot Loading…
[MGPU] Split test_broadcast_major_dim into two tests
#33892 opened Dec 11, 2025 by copybara-service bot Loading…
[pallas:mosaic] Allowed specifying tiling for SC_*_SUBCORE kernels
#33891 opened Dec 11, 2025 by copybara-service bot Loading…
@superbobry
[no-thunks] Implement FlatTree, an internal representation of pytrees. kokoro:force-run pull ready Ready for copybara import and testing
#33890 opened Dec 11, 2025 by dougalm Loading…
9
[Mosaic] NFC: Move common parts in tpu.td to a separate file.
#33888 opened Dec 11, 2025 by copybara-service bot Loading…
Follow up pr after triton integration cl and tokamax pr.
#33887 opened Dec 11, 2025 by copybara-service bot Loading…
[export] Fix the "with mesh" deprecation warning kokoro:force-run pull ready Ready for copybara import and testing
#33885 opened Dec 11, 2025 by gnecula Loading…
@gnecula
2
Complete reorg of tutorials and guides
#33884 opened Dec 11, 2025 by melissawm Loading…
@emilyfertig
4
Expand and update XProf docs
#33881 opened Dec 11, 2025 by mtsokol Loading…
12
[Pallas:MGPU] Allow repeated indexing of barriers
#33880 opened Dec 11, 2025 by copybara-service bot Loading…
@apaszke
Migrate TSAN workflow to use RBE.
#33867 opened Dec 10, 2025 by copybara-service bot Loading…
@belitskiy
[Pallas:MGPU] Support not tiled 3D transposed in swap LANE lowering rule.
#33864 opened Dec 10, 2025 by copybara-service bot Loading…
[pallas] Deprecated pltpu.ANY in favor of pl.ANY
#33861 opened Dec 10, 2025 by copybara-service bot Loading…
@superbobry
Disable failing TPU presubmit test.
#33857 opened Dec 10, 2025 by copybara-service bot Loading…
@danielsuo
[Mosaic GPU] Add support for indexing untiled dims
#33856 opened Dec 10, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add limited support for reshapes of tiled layouts
#33855 opened Dec 10, 2025 by copybara-service bot Loading…
@apaszke
[export] Add backwards compatibility tests for v6. kokoro:force-run pull ready Ready for copybara import and testing
#33853 opened Dec 10, 2025 by gnecula Loading…
@gnecula
3
Internal change to JAX lib visibility.
#33845 opened Dec 10, 2025 by copybara-service bot Loading…
Automated Code Change
#33843 opened Dec 10, 2025 by copybara-service bot Loading…
ProTip! Updated in the last three days: updated:>2025-12-08.