Skip to content

Pull requests: jax-ml/jax

New pull request New
Clear current search query, filters, and sorts
1 Open 75 Closed
1 Open 75 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 argument donation for xmap cla: yes pull ready Ready for copybara import and testing
#5632 by apaszke was merged Feb 8, 2021 Loading…
@apaszke
Implement a trivial ppermute collective batcher cla: yes pull ready Ready for copybara import and testing
#5557 by apaszke was merged Feb 1, 2021 Loading…
@apaszke
Add axis_sizes to xmap cla: yes pull ready Ready for copybara import and testing
#5602 by apaszke was merged Feb 2, 2021 Loading…
@apaszke
Substitute axis names in nested jaxprs cla: yes pull ready Ready for copybara import and testing
#5977 by apaszke was merged Mar 9, 2021 Loading…
@apaszke
Hotfix for psum transpose cla: yes pull ready Ready for copybara import and testing
#5648 by apaszke was merged Feb 5, 2021 Loading…
@apaszke
Fix a bug in xmap mesh slicing code cla: yes pull ready Ready for copybara import and testing
#5643 by apaszke was merged Feb 5, 2021 Loading…
@apaszke
Fix cache misses when re-creating equivalent mesh objects cla: yes
#6492 by copybara-service bot was merged Apr 20, 2021 Loading…
@apaszke
Add support for named axes in jax.nn.one_hot cla: yes pull ready Ready for copybara import and testing
#5600 by apaszke was merged Feb 2, 2021 Loading…
@apaszke
Always run xmap over all mesh axes cla: yes pull ready Ready for copybara import and testing
#6066 by apaszke was merged Jun 8, 2021 Loading…
@apaszke
Fix incorrect handling of constants in xmap's partial_eval rule cla: yes pull ready Ready for copybara import and testing
#6065 by apaszke was merged Mar 19, 2021 Loading…
@apaszke
Add a type checking rule for xmap cla: yes
#6505 by copybara-service bot was merged Apr 20, 2021 Loading…
@apaszke
Move vtile to batching.py, make it possible to add new BatchTraces cla: yes
#6504 by copybara-service bot was merged Apr 20, 2021 Loading…
@apaszke
Extend pjit error checking to rank errors too cla: yes
#6508 by copybara-service bot was merged Apr 20, 2021 Loading…
@apaszke
Small updates to abstract eval rules (AWN related) cla: yes pull ready Ready for copybara import and testing
#6419 by apaszke was merged Apr 15, 2021 Loading…
@apaszke
Insert xmap SPMD axes into pjit sharding annotations cla: yes
#6669 by copybara-service bot was merged Jun 3, 2021 Loading…
@apaszke
Fix incorrect handling of axis_index_groups in parallel primitive fallbacks cla: yes
#6818 by copybara-service bot was merged Jun 2, 2021 Loading…
@apaszke
Add new lowering APIs for xmap cla: yes
#7873 by copybara-service bot was closed Sep 10, 2021 Loading…
@apaszke
Remove collectives from all_gather transpose rule cla: yes pull ready Ready for copybara import and testing
#8115 by sharadmv was merged Oct 27, 2021 Loading…
@sharadmv @apaszke
Disallow jax2tf translation of multi-process pjits cla: yes pull ready Ready for copybara import and testing
#8154 by apaszke was merged Oct 11, 2021 Loading…
@apaszke
[Mosaic GPU] Use gettatr to import version-specific dialect ops pull ready Ready for copybara import and testing
#26523 by andportnoy was merged Feb 14, 2025 Loading…
@apaszke
Update GPU and NVGPU MLIR bindings to match upstream MLIR changes pull ready Ready for copybara import and testing
#20911 by apaszke was merged Apr 25, 2024 Loading…
@apaszke
[Pallas:MGPU] Add a diagram outlining GPU memory spaces documentation kokoro:force-run pull ready Ready for copybara import and testing
#28056 by apaszke was merged Apr 16, 2025 Loading…
@apaszke
Don't require ellipsis at the end of each list axis spec in xmap cla: yes pull ready Ready for copybara import and testing
#5633 by apaszke was merged Feb 5, 2021 Loading…
@apaszke
1
Add support for named axes in jnp.mean and jnp.std cla: yes pull ready Ready for copybara import and testing
#5644 by apaszke was merged Feb 9, 2021 Loading…
@apaszke
1
Clean up xmap tests cla: yes pull ready Ready for copybara import and testing
#5435 by apaszke was merged Jan 18, 2021 Loading…
@apaszke
1
ProTip! no:milestone will show everything without a milestone.