Skip to content

Pull requests: jax-ml/jax

New pull request New
536 Open 22,388 Closed
536 Open 22,388 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 ExecuteDeviceKernel frontend in JAX
#32869 opened Oct 25, 2025 by serach24 Loading…
5
[Pallas] Add output fusion support for reduce_sum/broadcast
#32868 opened Oct 25, 2025 by copybara-service bot Loading…
@sharadmv
Disable incompatible paged attention tests for TPU v7
#32867 opened Oct 24, 2025 by copybara-service bot Loading…
@justinjfu
[jax] Migrate jaxlib to xla::Future
#32866 opened Oct 24, 2025 by copybara-service bot Loading…
@ezhulenev
Tag some more TF-using tests as thread unsafe.
#32864 opened Oct 24, 2025 by copybara-service bot Loading…
@hawkinsp
Disable failing flaky tests on TAP.
#32863 opened Oct 24, 2025 by copybara-service bot Loading…
Temporarily disable XNNPACK for qdwh_test.
#32862 opened Oct 24, 2025 by copybara-service bot Loading…
Allow jax.jit to be used as a decorator factory.
#32861 opened Oct 24, 2025 by copybara-service bot Loading…
@jakevdp
[test] handle accelerated dep of lax.dot positional args
#32859 opened Oct 24, 2025 by jakevdp Loading…
@jakevdp
4
Set the concrete mesh to empty when executing a default callback
#32858 opened Oct 24, 2025 by copybara-service bot Loading…
@yashk2810
[MOSAIC] Do not use shape invariant numerics by default.
#32857 opened Oct 24, 2025 by copybara-service bot Loading…
Remove internal dependency and open-source multiprocess/array_serialization_test.
#32856 opened Oct 24, 2025 by copybara-service bot Loading…
@emilyfertig
Increase heartbeat timeout in multiprocess tests.
#32852 opened Oct 23, 2025 by copybara-service bot Loading…
@mwhittaker
[Mosaic] Fold no-op BitcastVregOp.
#32848 opened Oct 23, 2025 by copybara-service bot Loading…
Propagate effects in the abstract eval rule for custom_vmap_p.
#32847 opened Oct 23, 2025 by copybara-service bot Loading…
@jburnim
Change --jax_device_under_test to use "cuda" or "rocm".
#32845 opened Oct 23, 2025 by copybara-service bot Loading…
@hawkinsp
Fix use of deprecated jax_platform_name kokoro:force-run pull ready Ready for copybara import and testing
#32840 opened Oct 23, 2025 by Prince213 Loading…
6
[pallas:triton] Deprecated pl.max_contiguous -- it should now be accessed via the triton submodule
#32836 opened Oct 23, 2025 by copybara-service bot Loading…
@superbobry
Reverts 671483f4d9e50c38a168e07977aa35d0edfbb48b
#32827 opened Oct 23, 2025 by copybara-service bot Loading…
[Mosaic TPU] Allow small sublane tilings in second/third minor transpose for unpacked dtype instead of forcing native tiling, which can be more efficient when the second minor dimension of input tensor is small.
#32823 opened Oct 22, 2025 by copybara-service bot Loading…
ruleml test
#32820 opened Oct 22, 2025 by copybara-service bot Loading…
@bhatuzdaname
[pallas:sc] Added SparseCore specific attributes to TpuInfo
#32819 opened Oct 22, 2025 by copybara-service bot Loading…
@superbobry
[Pallas TPU] Add a primitive stochastic_round
#32813 opened Oct 22, 2025 by copybara-service bot Loading…
[Mosaic] Canonicalize CmpIOp
#32807 opened Oct 22, 2025 by copybara-service bot Loading…
[Pallas:GPU] Add Warpgroup lowering for lax.slice_p.
#32797 opened Oct 22, 2025 by copybara-service bot Loading…
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.