Skip to content

Pull requests: jax-ml/jax

New pull request New
592 Open 19,810 Closed
592 Open 19,810 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

[Pallas:MGPU] Missing div lowering for LoweringSemantics.Lane and PrimitiveSemantics.Warp.
#29953 opened Jul 3, 2025 by copybara-service bot Loading…
[Mosaic GPU] Implement meet and join for replicated layouts.
#29952 opened Jul 3, 2025 by copybara-service bot Loading…
@bchetioui
#sdy Fix forward of making XLA C++ changes so we can fall back to GSPMD in JAX export if the loaded module was lowered for GSPMD.
#29951 opened Jul 3, 2025 by copybara-service bot Loading…
[Pallas] Improve the error message when writing a ref to a ref
#29950 opened Jul 3, 2025 by copybara-service bot Loading…
@apaszke
[XLA:CPU][nanort] Test that compilation pipeline propagates int4 packing information to executable.
#29947 opened Jul 3, 2025 by copybara-service bot Loading…
Read corresponding environment variables for setting config flags.
#29945 opened Jul 2, 2025 by copybara-service bot Loading…
@danielsuo
[Mosaic] Add support for dynamic filter values
#29938 opened Jul 2, 2025 by copybara-service bot Loading…
Avoid strong refs to tracers in DynamicJaxprTrace. kokoro:force-run pull ready Ready for copybara import and testing
#29937 opened Jul 2, 2025 by dougalm Loading…
[Mosaic] Add tpu.wait_indirect_dma
#29935 opened Jul 2, 2025 by copybara-service bot Loading…
[jax:pallas] Remove weakref_lru_cache from _trace_kernel_to_jaxpr.
#29934 opened Jul 2, 2025 by copybara-service bot Loading…
@danielsuo
Rename OS-specific extension from PJRT plugin target names
#29933 opened Jul 2, 2025 by copybara-service bot Loading…
Use qualified names in tracebacks.
#29932 opened Jul 2, 2025 by copybara-service bot Loading…
@hawkinsp
Precompute has_changed and will_change during pallas pipelines.
#29931 opened Jul 2, 2025 by copybara-service bot Loading…
@justinjfu
[Mosaic GPU] Unify evaluate_equation and simplify_equation into reduce_equation.
#29928 opened Jul 2, 2025 by copybara-service bot Loading…
@bchetioui
Remove pretty_printer Python code.
#29927 opened Jul 2, 2025 by copybara-service bot Loading…
@hawkinsp
[CI] Update CI builds to Ubuntu 22.04
#29926 opened Jul 2, 2025 by copybara-service bot Loading…
@MichaelHudgins
[Mosaic GPU][NFC] Import equations as eqns in layout_inference2.py.
#29925 opened Jul 2, 2025 by copybara-service bot Loading…
@bchetioui
[Mosaic GPU] Add logic to derive default layouts in an undetermined system.
#29923 opened Jul 2, 2025 by copybara-service bot Loading…
@bchetioui
[Mosaic GPU] Propagate layouts from consumers and producers in the new layout inference.
#29922 opened Jul 2, 2025 by copybara-service bot Loading…
@bchetioui
[pallas] Support reshape with trailing 1s in fuser.
#29920 opened Jul 2, 2025 by copybara-service bot Loading…
@chr1sj0nes
Improve rendering of function names in jaxpr equation profile.
#29918 opened Jul 2, 2025 by copybara-service bot Loading…
@hawkinsp
Include non-IFRT-PjRT arrays in JAX heap profiles.
#29917 opened Jul 2, 2025 by copybara-service bot Loading…
Compute non-zero flops for gather with FILL_OR_DROP.
#29914 opened Jul 2, 2025 by copybara-service bot Loading…
Readd a name to the MLIR lowering of core.closed_call.
#29913 opened Jul 1, 2025 by copybara-service bot Loading…
@hawkinsp
Allow async remote copy to take device id as a dict of axis names and indices.
#29911 opened Jul 1, 2025 by copybara-service bot Loading…
@IvyZX
ProTip! Add no:assignee to see everything that’s not assigned.