Skip to content

Pull requests: jax-ml/jax

New pull request New
490 Open 18,065 Closed
490 Open 18,065 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

[Mosaic TPU] FWD compatibility needs to keep previous version at least one month.
#27762 opened Apr 6, 2025 by copybara-service bot Loading…
[vmappable] fix trace context bugs kokoro:force-run pull ready Ready for copybara import and testing
#27759 opened Apr 6, 2025 by mattjj Loading…
Address previous FP8-related TODOs in jaxlib/XLA.
#27757 opened Apr 5, 2025 by copybara-service bot Loading…
The ml_dtype requirement in JAX was updated to version 0.5.0+ (on Mar 20, 2025) - [commit 4b7ead4](https://github.com/jax-ml/jax/commit/4b7ead4d02f866077f11dcfcca7507533a441bcc)
#27754 opened Apr 5, 2025 by copybara-service bot Loading…
Reverts e2f67e0ef1af19f4d32e02f6fb927502469b32c6
#27752 opened Apr 4, 2025 by copybara-service bot Loading…
roll back test-breaking change
#27751 opened Apr 4, 2025 by copybara-service bot Loading…
Set varying_axes_in_types to True
#27749 opened Apr 4, 2025 by copybara-service bot Loading…
@yashk2810
Move l_next (float32) division out of the flash attention kernel
#27746 opened Apr 4, 2025 by copybara-service bot Loading…
Create a PjRt protos dir.
#27744 opened Apr 4, 2025 by copybara-service bot Loading…
Fix jax.random.randint for 8-bit and 16-bit dtypes kokoro:force-run pull ready Ready for copybara import and testing
#27742 opened Apr 4, 2025 by jakevdp Loading…
@jakevdp
3
Test download time.
#27741 opened Apr 4, 2025 by copybara-service bot Loading…
[pallas:mgpu] Un-undo ref transforms, ie annihilate transforms.
#27739 opened Apr 4, 2025 by copybara-service bot Loading…
[Mosaic GPU] Refactor and generalize code in optimization_barrier.
#27736 opened Apr 4, 2025 by copybara-service bot Loading…
Apply forwarding in pjit linearization rule to avoid intermediate copies kokoro:force-run pull ready Ready for copybara import and testing
#27735 opened Apr 4, 2025 by dfm Loading…
@dfm
Add jax.fwd_and_bwd
#27733 opened Apr 4, 2025 by gspschmid Loading…
@mattjj
1
Migrate jaxlib to use a single common .so file for all C++ dependencies. kokoro:force-run pull ready Ready for copybara import and testing windows:force-run Apply this label to a PR to force a run of the windows CI
#27732 opened Apr 4, 2025 by hawkinsp Loading…
1
[mutable-arrays] limit implicit ref_swap dtype promotion
#27728 opened Apr 4, 2025 by copybara-service bot Loading…
1
@mattjj
Plumbing for has-a between LoadedExec and Exec
#27727 opened Apr 4, 2025 by copybara-service bot Loading…
[JAX] Use xla::ifrt::Client::MakeArraysFromHostBufferShards() in Array creation when possible
#27722 opened Apr 4, 2025 by copybara-service bot Loading…
Fix shard_map for primitives with multiple outputs and None specs kokoro:force-run pull ready Ready for copybara import and testing
#27720 opened Apr 3, 2025 by justinjfu Loading…
13
jax.numpy: support __jax_array__ in remaining APIs kokoro:force-run pull ready Ready for copybara import and testing
#27716 opened Apr 3, 2025 by jakevdp Loading…
@jakevdp
Swap absl logging for Python logging
#27715 opened Apr 3, 2025 by ZacCranko • Draft
@ZacCranko
[Mosaic GPU] Add scaffolding for a new lowering "axis" (UserThreadSemantics), in addition to the existing ThreadSemantics (renamed to LoweringSemantics).
#27713 opened Apr 3, 2025 by copybara-service bot Loading…
@justinjfu
Move DropVar annotation into pretty printer pull ready Ready for copybara import and testing
#27709 opened Apr 3, 2025 by dfm Loading…
@dfm
Relax jax dependency constraints to be able to install RC wheels
#27704 opened Apr 3, 2025 by copybara-service bot Loading…
@nitins17
ProTip! Adding no:label will show everything without a label.