Skip to content

Pull requests: jax-ml/jax

New pull request New
797 Open 24,099 Closed
797 Open 24,099 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

[hijax] test vmap + hijax types + hijax primitive kokoro:force-run pull ready Ready for copybara import and testing
#35136 opened Feb 17, 2026 by mattjj Loading…
@mattjj
1
Reverts 4a14e0c2cce9ff41cbca271b429d36df8215e608
#35135 opened Feb 17, 2026 by copybara-service bot Loading…
Make small tweaks to fault tolerant mcjax blog.
#35133 opened Feb 17, 2026 by mwhittaker Loading…
@mwhittaker
1
Fix wheel_command to use wheel_build_command_base strings
#35132 opened Feb 17, 2026 by copybara-service bot Loading…
@ecalubaquib
Restore scan length check
#35131 opened Feb 17, 2026 by jakevdp Loading…
@jakevdp
1
[mosaic] Add optional core_type parameter to tpu.memory_space
#35127 opened Feb 17, 2026 by copybara-service bot Loading…
@superbobry
Add support for SDPA backward kernel with optimized flash attention kernel.
#35117 opened Feb 17, 2026 by copybara-service bot Loading…
[MGPU] Support scalar (int, float, bool) arguments in as_torch_kernel via static_argnums
#35116 opened Feb 16, 2026 by oulgen Loading…
@apaszke
5
[ROCm] Add hip_threefry2x32_ffi to stable custom call targets kokoro:force-run pull ready Ready for copybara import and testing
#35115 opened Feb 16, 2026 by AratiGanesh Loading…
3
[Pallas][Mosaic GPU] Add MGPU splash attention implementation
#35114 opened Feb 16, 2026 by justinjfu Loading…
6
[pallas] mpmd_map now supports lowering multiple kernels
#35113 opened Feb 16, 2026 by copybara-service bot Loading…
@superbobry
[pallas] Removed unnecessary api.named_scope from _pallas_call
#35112 opened Feb 16, 2026 by copybara-service bot Loading…
@superbobry
[ROCm] Add ROCm support for eigh export backwards compatibility tests
#35111 opened Feb 16, 2026 by AratiGanesh Loading…
@jakevdp
5
Add ormqr primitive and jax.scipy.linalg.qr_multiply
#35104 opened Feb 16, 2026 by jpbrodrick89 Loading…
5
[XLA:GPU] Add clique requirements which allows thunks to request additional barriers.
#35103 opened Feb 16, 2026 by copybara-service bot Loading…
[ROCm] Set release rpaths to rocm so targets
#35102 opened Feb 16, 2026 by alekstheod Loading…
5
Move the allow-hlo-dialects-in-mosaic flag to jax.config to fix a flag leak in tests when tpu_custom_call.py is imported.
#35097 opened Feb 14, 2026 by copybara-service bot Loading…
@rdyro
Fix variance_scaling division by zero error.
#35096 opened Feb 14, 2026 by carlosgmartin Loading…
1
Add documentation to lax.reduce and lax.reduce_window stating that init_values and init_value must consist of scalars kokoro:force-run pull ready Ready for copybara import and testing
#35095 opened Feb 14, 2026 by carlosgmartin Loading…
@jakevdp
4
Move core indexing functionality to _src/indexing.py
#35092 opened Feb 14, 2026 by jakevdp • Draft
@jakevdp
1
fixed scan crash on GPU with non-default matmul precision
#35091 opened Feb 14, 2026 by sarayusapa Loading…
4
fix an infinite recursion during lowering kokoro:force-run pull ready Ready for copybara import and testing
#35086 opened Feb 13, 2026 by v0i0 Loading…
4
Add a monitoring metric to track the duration of the JAX lowering process.
#35085 opened Feb 13, 2026 by copybara-service bot Loading…
Fix PRNGKeyArray pickling on 3.11.
#35079 opened Feb 13, 2026 by copybara-service bot Loading…
@belitskiy
Reverts 8ca2421b09e9b3310a3af6466d8eba4fd48999a3
#35078 opened Feb 13, 2026 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.