Skip to content

Pull requests: jax-ml/jax

New pull request New
649 Open 23,283 Closed
649 Open 23,283 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

Skip async store reduction tests for int64/uint64 when x64 is disabled.
#34002 opened Dec 17, 2025 by copybara-service bot Loading…
@belitskiy
[MGPU] Doc fix and shape size fix for broadcast in WGSplatFragLayout.
#33998 opened Dec 17, 2025 by copybara-service bot Loading…
Automated Code Change
#33997 opened Dec 17, 2025 by copybara-service bot Loading…
Reverts 6940903c006b9ea57af6e0043b4f86d4c9487045
#33995 opened Dec 17, 2025 by copybara-service bot Loading…
Automated Code Change
#33992 opened Dec 17, 2025 by copybara-service bot Loading…
adding few line of code ,this is for practice purpose so please ignore.
#33989 opened Dec 17, 2025 by Prasun-13 Loading…
4
[XLA:TPU] Abort early if the peak HLO memory consumption is greater than HBM
#33986 opened Dec 17, 2025 by copybara-service bot Loading…
jax.numpy: return tuples to match NumPy 2.0
#33984 opened Dec 16, 2025 by jakevdp Loading…
@jakevdp
3
[Pallas MGPU] Simplify how we keep track of the current output slices. Keeping track of the full store slices is unnecessary because the slice size doesn’t change, we really only care about the the start of the slices.
#33982 opened Dec 16, 2025 by copybara-service bot Loading…
@Rifur13
Bisect failing callback test
#33979 opened Dec 16, 2025 by copybara-service bot Loading…
@belitskiy
Deterministic ordering for sets when dumping jaxpr.
#33976 opened Dec 16, 2025 by copybara-service bot Loading…
internal changes
#33975 opened Dec 16, 2025 by copybara-service bot Loading…
[Bug] lax.padtype_to_pads: strip NumPy scalars from padding tuples
#33974 opened Dec 16, 2025 by Prakharprasun Loading…
1
[Pallas:MGPU] Add support for indexing untiled dimensions under WG semantic.
#33973 opened Dec 16, 2025 by copybara-service bot Loading…
Implement higher order interpolations in jax.scipy.ndimage
#33970 opened Dec 16, 2025 by hexane360 Loading…
1
8
Add state discharge rule for closed call/call primitives
#33968 opened Dec 16, 2025 by copybara-service bot Loading…
@sharadmv
Add platform name to xla::ifrt::Device
#33966 opened Dec 16, 2025 by copybara-service bot Loading…
v6e-8 diagnosing run
#33964 opened Dec 15, 2025 by copybara-service bot Loading…
@belitskiy
Add batching rule for with_memory_space_constraint primitive.
#33963 opened Dec 15, 2025 by copybara-service bot Loading…
[CI] Modify TPU v7x jobs to include bazel, and exclude some python versions
#33960 opened Dec 15, 2025 by copybara-service bot Loading…
@MichaelHudgins
Bump tj-actions/changed-files from 46.0.5 to 47.0.1 dependencies Pull requests that update a dependency file github_actions Pull requests that update Github_actions code
#33958 opened Dec 15, 2025 by dependabot bot Loading…
Bump medyagh/setup-minikube from 0.0.19 to 0.0.21 dependencies Pull requests that update a dependency file github_actions Pull requests that update Github_actions code
#33957 opened Dec 15, 2025 by dependabot bot Loading…
[PJRT] Simplify and broaden the interface of TransposePlan.
#33954 opened Dec 15, 2025 by copybara-service bot Loading…
@hawkinsp
Enable xla_gpu_experimental_scaled_dot_with_triton by default.
#33952 opened Dec 15, 2025 by copybara-service bot Loading…
Automated Code Change
#33946 opened Dec 15, 2025 by copybara-service bot Loading…
ProTip! Adding no:label will show everything without a label.