Skip to content

Pull requests: jax-ml/jax

New pull request New
556 Open 19,319 Closed
556 Open 19,319 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][Easy] Terser printing of GridMapping unless debug is set.
#29318 opened Jun 8, 2025 by copybara-service bot Loading…
[JAX] Add vma to ShapeDtypeStruct constructor arguments.
#29314 opened Jun 7, 2025 by copybara-service bot Loading…
@yashk2810
Clarify argument order for lax.associative_scan when reverse=True. kokoro:force-run pull ready Ready for copybara import and testing
#29313 opened Jun 7, 2025 by carlosgmartin Loading…
1
[Mosaic GPU] Make the Pallas Blackwell matmul kernel persistent.
#29312 opened Jun 6, 2025 by copybara-service bot Loading…
@Rifur13
Add numerics check for unreduced e2e.
#29310 opened Jun 6, 2025 by copybara-service bot Loading…
Port pretty-printer to C++.
#29308 opened Jun 6, 2025 by copybara-service bot Loading…
@hawkinsp
Adding jax._src.util.cache_clearing_funs to jax.extend.backend.
#29303 opened Jun 6, 2025 by copybara-service bot Loading…
Add alternative location of CUDA_ROOT for Bazel build/tests with hermetic CUDA.
#29300 opened Jun 6, 2025 by copybara-service bot Loading…
[CI] Run Mosaic H100 and B200 tests on all PRs that target mosaic subpaths
#29298 opened Jun 6, 2025 by MichaelHudgins • Draft
Pallas documentation fixes.
#29297 opened Jun 6, 2025 by copybara-service bot Loading…
//tests:scaled_matmul_stablehlo_test: fix for xla#27096 kokoro:force-run pull ready Ready for copybara import and testing
#29294 opened Jun 6, 2025 by olupton Loading…
[Mosaic GPU] Simplify how we skip transforms for unrealized casts and gpu shared memory.
#29293 opened Jun 6, 2025 by copybara-service bot Loading…
Delete extraneous block_until_readys on JAX dispatch benchmarks.
#29289 opened Jun 5, 2025 by copybara-service bot Loading…
Add test and refactor Device Assignment.
#29288 opened Jun 5, 2025 by copybara-service bot Loading…
Improve device assignment string format and add AbslStringify implementation.
#29287 opened Jun 5, 2025 by copybara-service bot Loading…
Implemented cross-host memory transfer on GPU.
#29286 opened Jun 5, 2025 by copybara-service bot Loading…
@mwhittaker
Add unit tests for parameter and optimizer state offload
#29284 opened Jun 5, 2025 by zhenying-liu Loading…
11
Expose local/global ExchangeTopologies timeouts for PJRT CPU client.
#29283 opened Jun 5, 2025 by copybara-service bot Loading…
Remove forward_compat as it is past the support date.
#29282 opened Jun 5, 2025 by copybara-service bot Loading…
[ROCm] ROCm7 Plugin Updates
#29281 opened Jun 5, 2025 by mrodden Loading…
1
#sdy don't make JAX export use SdyRoundTripExportPipeline to stringify attributes and convert ops to StableHLO CustomCallOps and back.
#29272 opened Jun 5, 2025 by copybara-service bot Loading…
[pallas:mgpu] Optionally skip allocating registers in warp specialized pipelining.
#29271 opened Jun 5, 2025 by copybara-service bot Loading…
[Mosaic GPU] Handle None transforms in swizzle_and_transforms_from_transforms_attr
#29269 opened Jun 5, 2025 by copybara-service bot Loading…
Change use_shardy_partitioner default to None instead of False.
#29262 opened Jun 5, 2025 by copybara-service bot Loading…
Fixed NaN propagation in optimal_step_size using nanmin/nanmax
#29261 opened Jun 5, 2025 by SuriyaaMM Loading…
ProTip! Updated in the last three days: updated:>2025-06-05.
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy