Skip to content

Pull requests: jax-ml/jax

New pull request New
490 Open 18,784 Closed
490 Open 18,784 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
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Add out_sharding to jnp.repeat. Drop into auto mode if out_sharding is provided.
#28600 opened May 8, 2025 by copybara-service bot Loading…
1
@yashk2810
[Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM) CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28595 opened May 7, 2025 by andportnoy Loading…
Use the C4A machine type for Linux Arm64 builds
#28591 opened May 7, 2025 by copybara-service bot Loading…
@nitins17
Record event start time in dispatch.LogElapsedTimeContextManager.
#28590 opened May 7, 2025 by copybara-service bot Loading…
@danielsuo
Consolidate initial/final style custom_vjp primitives into one pull ready Ready for copybara import and testing
#28589 opened May 7, 2025 by dfm Loading…
@dfm
[xla:ifrt] Remove references to xla::ifrt::Compile that return LoadedExecutables.
#28581 opened May 7, 2025 by copybara-service bot Loading…
@danielsuo
[Mosaic GPU] Use explicit load/store methods instead of __getitem__/__setitem__
#28580 opened May 7, 2025 by copybara-service bot Loading…
@apaszke
[pallas:mosaic_gpu] Removed the GPU* prefix from Mosaic GPU-specific types
#28578 opened May 7, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic] Fork transpose operation to TPU dialect.
#28569 opened May 7, 2025 by copybara-service bot Loading…
Fix random.binomial(...) to accept dtype=int
#28568 opened May 7, 2025 by aghilann Loading…
1
@jakevdp
8
Replace std::shared_ptr<xla::ifrt::LoadedExecutable> with xla::ifrt::LoadedExecutableRef
#28565 opened May 7, 2025 by copybara-service bot Loading…
Expose profiler_data submodule from XLA to Jaxlib.
#28564 opened May 6, 2025 by copybara-service bot Loading…
@ZacCranko
Fixed OOM on pytest cuda.
#28563 opened May 6, 2025 by mwhittaker • Draft
@mwhittaker
Trying to fix broken nvshmem build.
#28562 opened May 6, 2025 by mwhittaker • Draft
@mwhittaker
fix box bug in scan transpose kokoro:force-run pull ready Ready for copybara import and testing
#28561 opened May 6, 2025 by mattjj Loading…
@mattjj
Fixing broken fused_attention_stablehlo_test_gpu.
#28559 opened May 6, 2025 by mwhittaker • Draft
@mwhittaker
Fix empty string handling for cloud_tpu_cluster
#28554 opened May 6, 2025 by copybara-service bot Loading…
[ragged-paged-attn] Implement static kv cache quantization. (The scale of kv cache is a scalar float value)
#28553 opened May 6, 2025 by copybara-service bot Loading…
Fix typo in persistent_compilation_cache.md
#28549 opened May 6, 2025 by jjyyxx Loading…
Add collective_axes to run_scoped
#28546 opened May 6, 2025 by copybara-service bot Loading…
@apaszke
Improve PackSubelementsOp::verify
#28536 opened May 6, 2025 by copybara-service bot Loading…
[Mosaic] Use native bf16 ops for tanh, exp and log on TPUv6+.
#28531 opened May 6, 2025 by copybara-service bot Loading…
Pass RULES_PYTHON_REPO_DEBUG value to the docker container env
#28530 opened May 5, 2025 by copybara-service bot Loading…
[Pallas][Mosaic GPU] Add transpose support to tcgen05_mma CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28528 opened May 5, 2025 by copybara-service bot Loading…
@justinjfu
[Pallas][Mosaic GPU] Add collective support to Blackwell/tcgen05 MMA. CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28524 opened May 5, 2025 by copybara-service bot Loading…
@justinjfu
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.