Skip to content

Pull requests: jax-ml/jax

New pull request New
489 Open 18,809 Closed
489 Open 18,809 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 host offloading docs to public website documentation
#28632 opened May 9, 2025 by copybara-service bot Loading…
@yashk2810
Annotate lower triangular matrix
#28631 opened May 8, 2025 by copybara-service bot Loading…
Ensure __jax_array__ support in binary ops kokoro:force-run pull ready Ready for copybara import and testing
#28630 opened May 8, 2025 by jakevdp Loading…
@jakevdp
[Pallas/TPU]
#28628 opened May 8, 2025 by copybara-service bot Loading…
@sharadmv
[ifrt] Refactor away from deprecated constructors
#28627 opened May 8, 2025 by copybara-service bot Loading…
Remove __div__ and __rdiv__ from jax.Array pull ready Ready for copybara import and testing
#28626 opened May 8, 2025 by jakevdp Loading…
@jakevdp
[Mosaic] Move sitofp lowering to Mosaic.
#28625 opened May 8, 2025 by copybara-service bot Loading…
Bumping oldest libtpu version for 12 week window.
#28621 opened May 8, 2025 by mwhittaker • Draft
@mwhittaker
Propagate use_shardy_partitioner to XlaCallModule op.
#28616 opened May 8, 2025 by copybara-service bot Loading…
[JAX:ANN] Make a test more reliable by increasing the sample size
#28615 opened May 8, 2025 by copybara-service bot Loading…
Add a pretty printing rule for custom_jvp
#28608 opened May 8, 2025 by dfm Loading…
@dfm
[platform_dependent] Ensure that platform_dependent only lowers for intended platforms
#28607 opened May 8, 2025 by gnecula Loading…
1
@gnecula
Define a type alias for tsl::RCReference<Value>
#28604 opened May 8, 2025 by copybara-service bot Loading…
[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
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…
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
ProTip! Mix and match filters to narrow down what you’re looking for.