Skip to content

Pull requests: jax-ml/jax

New pull request New
652 Open 23,267 Closed
652 Open 23,267 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 reshard to shard_map under full explicit mode if the input aval's sharding differs from in_specs passed to shard_map. Since in_specs of shard_map signal a reshard historically, we follow the same semantics here. Without this, the primal and cotangent types don't match.
#33988 opened Dec 17, 2025 by copybara-service bot Loading…
@yashk2810
[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
Make c-api topology and PjRtClient versions produce identical platform_verison
#33983 opened Dec 16, 2025 by copybara-service bot Loading…
@pschuh
[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
Skip test_scalar_debug_check in tpu_sparsecore_pallas_debug_check_test.
#33981 opened Dec 16, 2025 by copybara-service bot Loading…
@belitskiy
Bisect failing callback test
#33979 opened Dec 16, 2025 by copybara-service bot Loading…
@belitskiy
Reverts 58fbb8a3b7d4ecb860c4cf678cae1dd6f1079846
#33978 opened Dec 16, 2025 by copybara-service bot Loading…
@gnecula
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…
Reverts a99f76f6a6556223d71a2a63cd03019634085cc5
#33940 opened Dec 13, 2025 by copybara-service bot Loading…
Fix vjp type annotation
#33938 opened Dec 13, 2025 by copybara-service bot Loading…
@yashk2810
ProTip! no:milestone will show everything without a milestone.