Skip to content

Pull requests: jax-ml/jax

New pull request New
800 Open 23,985 Closed
800 Open 23,985 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

Request a full h100 gpu when neccessary.
#35014 opened Feb 12, 2026 by copybara-service bot Loading…
[ROCm] Remove incorrect ROCm lowering for scaled_matmul to prevent segfault
#35011 opened Feb 12, 2026 by Ruturaj4 Loading…
3
[typing] fix some pyrefly unbound local errors
#35010 opened Feb 12, 2026 by jakevdp Loading…
@jakevdp
1
Deprecate Pallas pltpu.repeat
#35009 opened Feb 11, 2026 by copybara-service bot Loading…
@sharadmv
[pmap] Accelerate deprecation of PmapSharding.
#35007 opened Feb 11, 2026 by copybara-service bot Loading…
@danielsuo
[Mosaic GPU] Implement tile_shape, a MLIR operation that allows tiling a memref.
#35006 opened Feb 11, 2026 by copybara-service bot Loading…
@bchetioui
Add a light, blocking Mosaic B200 presubmit.
#35005 opened Feb 11, 2026 by copybara-service bot Loading…
@belitskiy
[typing] fix pyrefly errors in jax/_src/numpy
#35004 opened Feb 11, 2026 by jakevdp Loading…
@jakevdp
4
Reimplement weakref_lru_cache on top of a custom hash map type.
#35003 opened Feb 11, 2026 by copybara-service bot Loading…
@hawkinsp
testing for update
#35001 opened Feb 11, 2026 by copybara-service bot Loading…
@ecalubaquib
[mpmd] Handle sharding-in-types in MPMD primitive lowerings.
#34999 opened Feb 11, 2026 by copybara-service bot Loading…
[XLA:MGPU] Port transfer plan classes to C++.
#34997 opened Feb 11, 2026 by copybara-service bot Loading…
[XLA:GPU] Remove unused collective metadata thunk.
#34996 opened Feb 11, 2026 by copybara-service bot Loading…
[Mosaic:GPU] Use collective kernel API directly.
#34995 opened Feb 11, 2026 by copybara-service bot Loading…
[shape_poly] Fix handling of shape polymorphism for pallas_call_batching kokoro:force-run pull ready Ready for copybara import and testing
#34988 opened Feb 11, 2026 by gnecula Loading…
@gnecula
6
docs: fix cookbook training loop math
#34987 opened Feb 11, 2026 by agarkovv Loading…
3
[pallas:sc] Slightly tweaked PipelineTest.test_basic under 7x
#34984 opened Feb 11, 2026 by copybara-service bot Loading…
@superbobry
[hijax] allow vmap/scan of hi types kokoro:force-run pull ready Ready for copybara import and testing
#34977 opened Feb 11, 2026 by mattjj Loading…
@mattjj
3
[typing] more pyrefly fixes kokoro:force-run pull ready Ready for copybara import and testing
#34975 opened Feb 10, 2026 by jakevdp Loading…
@jakevdp
21
[ROCm] Added ROCm tests to device_test.py unit test file.
#34974 opened Feb 10, 2026 by tsrw2048 Loading…
@gnecula
7
Add freethreaded python support in CUDA RBE presubmit.
#34972 opened Feb 10, 2026 by copybara-service bot Loading…
@ybaturina
[ROCm] fix the performance issue when n=1 or 2 AMD GPU Issues pertaining to AMD GPUs (ROCM) GPU Issues that apply to both CUDA and ROCM GPUs
#34971 opened Feb 10, 2026 by cj401-amd Loading…
3
[Mosaic TPU] Add support for bf16 sin, cos, and tan.
#34970 opened Feb 10, 2026 by copybara-service bot Loading…
Update test_cache_miss_explanations_new_function_in_loop to be robust to different backends
#34969 opened Feb 10, 2026 by copybara-service bot Loading…
[ROCm] Skip test_batch_axis_sharding_jvp
#34966 opened Feb 10, 2026 by AratiGanesh Loading…
3
ProTip! Updated in the last three days: updated:>2026-02-08.