Skip to content

Pull requests: jax-ml/jax

New pull request New
697 Open 23,492 Closed
697 Open 23,492 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

Revert switch to using cp.async.bulk.
#34296 opened Jan 13, 2026 by copybara-service bot Loading…
@chr1sj0nes
[Mosaic] Allow TC<->SC communication
#34295 opened Jan 13, 2026 by copybara-service bot Loading…
[dep] Deprecate jax.numpy.fix
#34293 opened Jan 12, 2026 by jakevdp Loading…
@jakevdp
3
[TPU:Mosaic] Optimize tpu.dynamic_rotate with tpu.assume <max>
#34292 opened Jan 12, 2026 by copybara-service bot Loading…
[TPU:Mosaic] Introduce tpu.assume op.
#34291 opened Jan 12, 2026 by copybara-service bot Loading…
[indexing] support newaxis in static/dynamic slice strategies
#34290 opened Jan 12, 2026 by jakevdp Loading…
@jakevdp
4
Update rules_ml_toolchain version to incorporate new folders structure.
#34289 opened Jan 12, 2026 by copybara-service bot Loading…
@ybaturina
1
Apply artificial shared memory limit to Pallas/Mosaic GPU tests
#34288 opened Jan 12, 2026 by olupton • Draft
4
[Mosaic GPU] Support more reduction kinds and layouts in MultiDimReductionOp lowering.
#34287 opened Jan 12, 2026 by copybara-service bot Loading…
[Pallas/interpreter] Add a prototype for a GPU kernel interpreter.
#34286 opened Jan 12, 2026 by copybara-service bot Loading…
[XLA:MGPU] Port warp_indices and lane_indices to C++.
#34285 opened Jan 12, 2026 by copybara-service bot Loading…
[XLA:MGPU] Port Partitioned.*Dims and VectorLength methods to C++.
#34279 opened Jan 12, 2026 by copybara-service bot Loading…
[Mosaic GPU] Lock down 8-bit pointwise ops (except bitwise and cmp) + add lshift/rshift
#34278 opened Jan 12, 2026 by copybara-service bot Loading…
@apaszke
Test return from constructor failure.
#34277 opened Jan 12, 2026 by copybara-service bot Loading…
Test XLA presubmits - Test port TiledLayout construction logic
#34275 opened Jan 12, 2026 by copybara-service bot Loading…
Automated Code Change
#34273 opened Jan 12, 2026 by copybara-service bot Loading…
[mosaic gpu] Add tests for TMA OOB masking with padded SMEM buffers kokoro:force-run pull ready Ready for copybara import and testing
#34272 opened Jan 12, 2026 by oulgen Loading…
5
Add lax.tile_p
#34270 opened Jan 11, 2026 by copybara-service bot Loading…
@levskaya
[Mosaic GPU][NFC] Replace duplicated logic with utils.get_contiguous_strides.
#34266 opened Jan 10, 2026 by copybara-service bot Loading…
Document cross-host device_put and running on a subset of devices. documentation
#34264 opened Jan 9, 2026 by emilyfertig Loading…
@emilyfertig
17
[XLA:TPU] Abort early if the peak HLO memory consumption is greater than HBM
#34260 opened Jan 9, 2026 by copybara-service bot Loading…
[doc] update dynamic_slice and related docstrings with clipping behavior documentation
#34257 opened Jan 9, 2026 by jakevdp Loading…
@jakevdp
1
Test Linux tests fix.
#34252 opened Jan 9, 2026 by copybara-service bot Loading…
@ybaturina
1
Set maxsize=None in trace_to_jaxpr weakref_lru_cache to preserve the same behavior as lu.cache
#34251 opened Jan 9, 2026 by copybara-service bot Loading…
@yashk2810
Test repro
#34248 opened Jan 9, 2026 by copybara-service bot Loading…
ProTip! Exclude everything labeled bug with -label:bug.