Skip to content

Pull requests: jax-ml/jax

New pull request New
494 Open 18,864 Closed
494 Open 18,864 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

[pallas:mosaic_gpu] Added support for unrolling to lax.fori_loop lowering
#28708 opened May 13, 2025 by copybara-service bot Loading…
@superbobry
Remove support for dimension level type.
#28705 opened May 13, 2025 by copybara-service bot Loading…
[Mosaic TPU][NFC] Consolidate getIntConst.
#28704 opened May 13, 2025 by copybara-service bot Loading…
Use DmaCopyChunk::Make because directly assigning the struct
#28703 opened May 13, 2025 by copybara-service bot Loading…
@pschuh
Infer sharding from arguments only if all mesh axes are explicit. Error out requiring in_specs/in_axes in all other cases.
#28702 opened May 13, 2025 by copybara-service bot Loading…
@yashk2810
Add use_raw_buffers which allows switching the implementation to
#28701 opened May 12, 2025 by copybara-service bot Loading…
@pschuh
[pallas] Pulled runtime_assert_enabled from pltpu to pl
#28699 opened May 12, 2025 by copybara-service bot Loading…
@superbobry
[Pallas] Allow f8 casting tests on TPUv5-.
#28698 opened May 12, 2025 by copybara-service bot Loading…
Improve the log for FlattenUpTo. Previous it's hard to know which key is missing when the dict is very large and got truncated in log.
#28697 opened May 12, 2025 by copybara-service bot Loading…
Bump fonttools from 4.51.0 to 4.58.0 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#28694 opened May 12, 2025 by dependabot bot Loading…
Bump hypothesis from 6.102.4 to 6.131.15 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#28693 opened May 12, 2025 by dependabot bot Loading…
@jakevdp
Deprecate the no-op custom_jvp_call_jaxpr_p import stub.
#28691 opened May 12, 2025 by copybara-service bot Loading…
@dfm
Revert pytest: use importlib mode by default kokoro:force-run pull ready Ready for copybara import and testing
#28690 opened May 12, 2025 by jakevdp Loading…
1
@jakevdp
2
[pallas:mgpu] Arrive only once to the pipeline input barriers.
#28689 opened May 12, 2025 by copybara-service bot Loading…
[CI] Add additional hardware to continuous non-rbe testing
#28688 opened May 12, 2025 by copybara-service bot Loading…
@MichaelHudgins
[pallas:mosaic_gpu] arrive_expect_tx primitive.
#28686 opened May 12, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add layout inference and lowering for scf.WhileOp and enable tests.
#28684 opened May 12, 2025 by copybara-service bot Loading…
[mosaic-gpu] add multicast ptr support to TMA with overlapped gemm and all reduce examples
#28679 opened May 12, 2025 by Amir-19 Loading…
[shard-map] start adding systematic smap tests kokoro:force-run pull ready Ready for copybara import and testing
#28665 opened May 10, 2025 by mattjj Loading…
@mattjj
[Mosaic] Make tpu.relayout an explicit operation, merge in existing behavior, stop calling relayout() in apply.
#28663 opened May 10, 2025 by copybara-service bot Loading…
JEP 28661: the __jax_array__ protocol
#28661 opened May 9, 2025 by jakevdp Loading…
@jakevdp
Fixed @@pypi//nvidia_nvshmem_cu12 errors.
#28658 opened May 9, 2025 by copybara-service bot Loading…
@mwhittaker
Bumping oldest libtpu version
#28652 opened May 9, 2025 by mwhittaker • Draft
@mwhittaker
Skip CSR matmat and matvec float tests on ROCm <6.4 (NaN issue with beta==0) kokoro:force-run pull ready Ready for copybara import and testing
#28635 opened May 9, 2025 by psanal35 Loading…
@dfm
5
Annotate lower triangular matrix
#28631 opened May 8, 2025 by copybara-service bot Loading…
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.