Skip to content

Pull requests: jax-ml/jax

New pull request New
564 Open 19,509 Closed
564 Open 19,509 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

Run pyupgrade --py311-plus.
#29550 opened Jun 18, 2025 by hawkinsp Loading…
Finalize a number of deprecations for JAX v0.7.0 kokoro:force-run pull ready Ready for copybara import and testing
#29549 opened Jun 18, 2025 by jakevdp Loading…
@jakevdp
2
Do not change the sharding type if there is only one devices.
#29548 opened Jun 18, 2025 by copybara-service bot Loading…
[Mosaic TPU] Make the backward-compatibility libtpu condition stricter
#29547 opened Jun 18, 2025 by copybara-service bot Loading…
@apaszke
Use bitwidth to check swizzle size instead of byte size.
#29546 opened Jun 18, 2025 by copybara-service bot Loading…
Reverts 366a7dfe5c5ec8f5f8b76460490dd13edf31c1e6
#29545 opened Jun 18, 2025 by copybara-service bot Loading…
@hawkinsp
Create a test suite for Pallas mosaic GPU tests.
#29542 opened Jun 18, 2025 by copybara-service bot Loading…
WIP: Refactor "How to think in JAX" and surrounding pages
#29541 opened Jun 18, 2025 by melissawm • Draft
Reenable AVX512 after LLVM fix upstream.
#29540 opened Jun 18, 2025 by copybara-service bot Loading…
[debug_info] Improve debug info for loop constructs. pull ready Ready for copybara import and testing
#29539 opened Jun 18, 2025 by gnecula Loading…
@gnecula
1
Pass shardy option through jax config.
#29536 opened Jun 18, 2025 by copybara-service bot Loading…
[Mosaic:TPU][Relayouts] Express some retilings as unpack+pack and remove offset restrictions
#29534 opened Jun 18, 2025 by copybara-service bot Loading…
Fix rare error with Literal in DynamicJaxprTracer.full_lower.
#29533 opened Jun 18, 2025 by copybara-service bot Loading…
@jburnim
Add select_n to roofline.
#29532 opened Jun 18, 2025 by copybara-service bot Loading…
[Pallas] Pass a kernel_info dict to mosaic
#29531 opened Jun 18, 2025 by copybara-service bot Loading…
@sharadmv
Remove the PjRt-IFRT dependency from the DCN transfer library.
#29527 opened Jun 17, 2025 by copybara-service bot Loading…
@emilyfertig
Remove deprecated LayoutUtil::IsDenseArray
#29525 opened Jun 17, 2025 by copybara-service bot Loading…
Skip pytest for Python 3.14 during the JAX release process
#29524 opened Jun 17, 2025 by copybara-service bot Loading…
Add logging for pjit infer_params cache hits and misses.
#29522 opened Jun 17, 2025 by copybara-service bot Loading…
@danielsuo
advertise plugin_attributes and consume it
#29517 opened Jun 17, 2025 by copybara-service bot Loading…
[mosaic] MemRef{Slice,Squeeze} verifiers now support strided layouts
#29514 opened Jun 17, 2025 by copybara-service bot Loading…
@superbobry
[doc] Clarify Profiling Docs for XProf and Tensorboard integration
#29509 opened Jun 17, 2025 by copybara-service bot Loading…
Add very basic support for shard_map + unreduced.
#29508 opened Jun 16, 2025 by copybara-service bot Loading…
@yashk2810
Use more concise tracer reprs by default.
#29507 opened Jun 16, 2025 by jakevdp Loading…
@jakevdp
2
Add an option to enable GPU collective cancelling.
#29505 opened Jun 16, 2025 by copybara-service bot Loading…
@mwhittaker
ProTip! Add no:assignee to see everything that’s not assigned.