Skip to content

Pull requests: jax-ml/jax

New pull request New
568 Open 22,776 Closed
568 Open 22,776 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

[Pallas:MGPU] Fix plgpu.inline_mgpu support for nested blocks.
#33337 opened Nov 14, 2025 by copybara-service bot Loading…
[Pallas:SC] Enable load_gather_masked tests with TC tiling
#33336 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:SC] Test with both SC and TC tiling
#33335 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[mpmd] Modify tests to also test custom pipeline API path
#33334 opened Nov 14, 2025 by copybara-service bot Loading…
[mpmd] Add custom pipeline API for user-defined scheduling and merging
#33333 opened Nov 14, 2025 by copybara-service bot Loading…
[mpmd] Move pipeline scheduling types and conversion functions to OSS
#33332 opened Nov 14, 2025 by copybara-service bot Loading…
[mpmd] Fix nanobind compatibility and add OPTIMIZE enum value
#33331 opened Nov 14, 2025 by copybara-service bot Loading…
Reduce the maximum test parallelism on Windows machines to avoid OOMs
#33330 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[DO NOT SUBMIT] Inject failure into CI
#33327 opened Nov 14, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Support 1D input and keepdims=False for argmax/argmin.
#33321 opened Nov 14, 2025 by copybara-service bot Loading…
[xla:cpu] Don't hoist small loops with runtime calls
#33319 opened Nov 14, 2025 by copybara-service bot Loading…
@ezhulenev
Add test of documentation coverage
#33316 opened Nov 13, 2025 by jakevdp Loading…
@jakevdp
5
[pallas fuser] transpose push rule returns block shape as tuple
#33315 opened Nov 13, 2025 by copybara-service bot Loading…
Add a vlog for the launch_id JAX chooses to aid debugging.
#33313 opened Nov 13, 2025 by copybara-service bot Loading…
@hawkinsp
If the matrix is square, do not assert if the gradient of the SVD is requested when compute_uv and full_matrices are both true.
#33312 opened Nov 13, 2025 by copybara-service bot Loading…
Set xla_tpu_enable_mlir_verifier in tests
#33311 opened Nov 13, 2025 by copybara-service bot Loading…
[Pallas] Skip TPU in pipelining scheduler tests.
#33308 opened Nov 13, 2025 by copybara-service bot Loading…
@justinjfu
[Mosaic GPU] Use a recursive tree implementation for vector_concat
#33305 opened Nov 13, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Tune the packing length for i4 -> f8e4m3fn conversion
#33304 opened Nov 13, 2025 by copybara-service bot Loading…
@apaszke
[mosaic] Made GetMosaicHardwareGeneration public
#33301 opened Nov 13, 2025 by copybara-service bot Loading…
@superbobry
Automated Code Change
#33293 opened Nov 13, 2025 by copybara-service bot Loading…
[pmap] Implement device_put_sharded using NamedSharding.
#33292 opened Nov 13, 2025 by copybara-service bot Loading…
@danielsuo
Enable RBE testing of JAX CUDA 13 wheels by loading CUDA UMD libs from Python code.
#33291 opened Nov 13, 2025 by copybara-service bot Loading…
@ybaturina
1
[pmap] Deprecate jax.sharding.PmapSharding, jax.device_put_replicated and jax.device_put_sharded.
#33290 opened Nov 13, 2025 by copybara-service bot Loading…
@danielsuo
[JAX] Re-enable grad(sin(x)) test case in xla_metadata_test.
#33287 opened Nov 12, 2025 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.