- Notifications
You must be signed in to change notification settings - Fork3.3k
Pull requests: jax-ml/jax
Author
Uh oh!
There was an error while loading.Please reload this page.
Label
Uh oh!
There was an error while loading.Please reload this page.
Projects
Uh oh!
There was an error while loading.Please reload this page.
Milestones
Uh oh!
There was an error while loading.Please reload this page.
Reviews
Assignee
Assigned to nobodyLoading
Uh oh!
There was an error while loading.Please reload this page.
Sort
Pull requests list
[Mosaic GPU] Implement GPU module unloading for Mosaic GPU custom calls.
#33599 openedNov 28, 2025 bycopybara-servicebotLoading…
[pallas:sc] Skip a few tests failing when the compiler uses tiled memrefs
#33598 openedNov 28, 2025 bycopybara-servicebotLoading…
[export] Add support for explicit sharding. pull readyReady for copybara import and testing
#33597 openedNov 28, 2025 bygneculaLoading…
Add programmatic profiling test with session_id support in JAX.
#33596 openedNov 28, 2025 bycopybara-servicebotLoading…
Add ppf (inverse CDF) for Gamma and Beta distributions (#20358)
#33589 openedNov 28, 2025 byVARUN3WARELoading…
[Mosaic GPU] Add support for arbitrary bitwidths in 32x32b TMEM transfers
#33580 openedNov 27, 2025 bycopybara-servicebotLoading…
[Mosaic] Support packed BroadcastInSublanesOp.
#33566 openedNov 26, 2025 bycopybara-servicebotLoading…
[Pallas/Triton, ROCm] Add lax.is_finite_p() lowering with tests kokoro:force-run pull readyReady for copybara import and testing
#33564 openedNov 26, 2025 byArech8Loading…
[Pallas/Mosaic GPU][NFC] Slightly clarify the synchronization requirements in the docstring of
copy_gmem_to_smem. #33563 openedNov 26, 2025 bycopybara-servicebotLoading…
Add exceptions for
numpyufuncs intests/lax_numpy_test.py::NumpySignaturesTest.testWrappedSignatureMatch. #33560 openedNov 26, 2025 bydanielsuoLoading…
Increase timeout for scaled_matmul_stablehlo_test. This is probably related to the increased autotuning time due to enabling TMA.
#33555 openedNov 26, 2025 bycopybara-servicebotLoading…
Fix 'Invalid Notebook' error in GitHub Preview of 'collab_cpu.ipynb'
#33551 openedNov 26, 2025 byashahbaLoading…
[xla:ffi] Move GPU FFI implementation to backends/gpu
#33549 openedNov 26, 2025 bycopybara-servicebotLoading…
[Pallas TPU] Add mesh axis info to pallas call metadata
#33546 openedNov 25, 2025 bycopybara-servicebotLoading…
[types] improve internal type annotations for indexing routines kokoro:force-run pull readyReady for copybara import and testing
#33544 openedNov 25, 2025 byjakevdpLoading…
[Pallas] Update instruction reference documentation kokoro:force-run pull readyReady for copybara import and testing
#33541 openedNov 25, 2025 byjustinjfuLoading…
Optimize lpmn_values using O(n²) Associated Legendre recurrences (#33370)
#33539 openedNov 25, 2025 byAKash-A007Loading…
ProTip! What’s not been updated in a month:updated:<2025-10-29.