-
Notifications
You must be signed in to change notification settings - Fork 3.2k
Pull requests: jax-ml/jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Transposed ragged dot kernel on Blackwell GPU
#33352
opened Nov 16, 2025 by
copybara-service
bot
Loading…
[Pallas/Fuser] Add basic for some simple output concatenation fusion
#33345
opened Nov 15, 2025 by
copybara-service
bot
Loading…
Test new rules_ml_toolchain version for CUDNN10.
#33339
opened Nov 14, 2025 by
copybara-service
bot
Loading…
[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…
[Pallas:SC] Test with both SC and TC tiling
#33335
opened Nov 14, 2025 by
copybara-service
bot
Loading…
[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…
[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…
Add test of documentation coverage
pull ready
Ready for copybara import and testing
#33316
opened Nov 13, 2025 by
jakevdp
Loading…
Set xla_tpu_enable_mlir_verifier in tests
#33311
opened Nov 13, 2025 by
copybara-service
bot
Loading…
[Mosaic GPU] Use a recursive tree implementation for vector_concat
#33305
opened Nov 13, 2025 by
copybara-service
bot
Loading…
[Mosaic GPU] Tune the packing length for i4 -> f8e4m3fn conversion
#33304
opened Nov 13, 2025 by
copybara-service
bot
Loading…
[mosaic] Made
GetMosaicHardwareGeneration public
#33301
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…
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…
Previous Next
ProTip!
Find all pull requests that aren't related to any open issues with -linked:issue.