Skip to content

Pull requests: jax-ml/jax

New pull request New
791 Open 23,903 Closed
791 Open 23,903 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

Automated Code Change
#34914 opened Feb 8, 2026 by copybara-service bot Loading…
Reverts e121d21893d46a613685363f1321dfe336ffbd2b
#34913 opened Feb 8, 2026 by copybara-service bot Loading…
[RFC] Enabled unused-variable in Ruff's config
#34911 opened Feb 7, 2026 by superbobry Loading…
66
[Take 2] In explicit mode, error out by default if the specs of arguments do not match in_specs passed to shard_map. In other words, in_specs in explicit_mode acts as an assert to make sure (jax.typeof(arg).sharding.spec == in_specs)
#34910 opened Feb 7, 2026 by copybara-service bot Loading…
@yashk2810
[jax:ffi] Do not use internal XLA ffi_api.h header
#34908 opened Feb 7, 2026 by ezhulenev Loading…
1
Reverts a7a4b6c7e1dceaf9ad49937170997dfa3a135c13
#34906 opened Feb 7, 2026 by copybara-service bot Loading…
Automated Code Change
#34905 opened Feb 7, 2026 by copybara-service bot Loading…
Temporary disable TMA with collective metadata is OSS due to linking issue
#34904 opened Feb 7, 2026 by copybara-service bot Loading…
Skipping testEighTinyNorm due to hipSolver issues
#34902 opened Feb 6, 2026 by AratiGanesh Loading…
3
Enable two test cases in jax/pjit_test.py.
#34898 opened Feb 6, 2026 by copybara-service bot Loading…
Integrate CCCL 3.2 headers in JAX.
#34897 opened Feb 6, 2026 by copybara-service bot Loading…
@ybaturina
Add ROCm backward compatibility test for cholesky solver
#34894 opened Feb 6, 2026 by AratiGanesh Loading…
3
[ROCm] Enable test deviceless aot compile test on ROCm
#34893 opened Feb 6, 2026 by magaonka-amd Loading…
3
Update rules_python to latest 1.8.3 for tensorflow and xla
#34892 opened Feb 6, 2026 by copybara-service bot Loading…
@ecalubaquib
Test out a suspect commit.
#34891 opened Feb 6, 2026 by copybara-service bot Loading…
@belitskiy
[ROCm] Enable memory space export tests on ROCm GPUs
#34890 opened Feb 6, 2026 by magaonka-amd Loading…
4
Update Flash Attention Kernel to support kernel export with dynamic shapes.
#34889 opened Feb 6, 2026 by copybara-service bot Loading…
Pass unused recoverable option to coordination service.
#34888 opened Feb 6, 2026 by copybara-service bot Loading…
@mwhittaker
Update rules_python to latest 1.8.3 for tensorflow and xla
#34887 opened Feb 6, 2026 by copybara-service bot Loading…
@ecalubaquib
[pallas:sc] Generalized supported vector shapes check to 7x
#34885 opened Feb 6, 2026 by copybara-service bot Loading…
@superbobry
Skip a multi-device test earlier if there's only one device
#34884 opened Feb 6, 2026 by olupton Loading…
3
[ROCm] Add ROCm backward compatibility test for threefry2x32
#34875 opened Feb 5, 2026 by AratiGanesh Loading…
3
[ROCm] Add ROCm backward compatibility test for lu_pivots_to_permutation
#34870 opened Feb 5, 2026 by AratiGanesh Loading…
4
[ROCm] Added back compatibility test for hipsolver_gesvd
#34869 opened Feb 5, 2026 by tsrw2048 Loading…
3
Simplified recoverable options in coordination service
#34868 opened Feb 5, 2026 by copybara-service bot Loading…
@mwhittaker
ProTip! Updated in the last three days: updated:>2026-02-05.