Skip to content

Pull requests: jax-ml/jax

New pull request New
583 Open 19,607 Closed
583 Open 19,607 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

Move jax._src.ops to jax._src.numpy build rule
#29695 opened Jun 24, 2025 by copybara-service bot Loading…
@jakevdp
Removed a few old jaxlib version guards
#29693 opened Jun 24, 2025 by copybara-service bot Loading…
@superbobry
[jax:pjit] Add jaxlib version guards to ArrayPjitTest.jit_mul_sum_sharding_preserved.
#29692 opened Jun 24, 2025 by copybara-service bot Loading…
@danielsuo
Support shardy in orbax.
#29691 opened Jun 24, 2025 by copybara-service bot Loading…
Allow all shardings if exported.nr_devices is 1 in _export.py.
#29690 opened Jun 24, 2025 by copybara-service bot Loading…
Custom batching rule for Cholesky that ensures symmetrize.
#29689 opened Jun 24, 2025 by copybara-service bot Loading…
[mosaic] Removed unused arguments from various tpu_custom_call functions
#29687 opened Jun 24, 2025 by copybara-service bot Loading…
@superbobry
[pallas:mosaic] Added device_id and core_id to tpu.dma_wait2 op
#29682 opened Jun 24, 2025 by copybara-service bot Loading…
@superbobry
Add JAX documentation for psend and precv API
#29677 opened Jun 24, 2025 by copybara-service bot Loading…
implements non-factorially scaled jet computation
#29676 opened Jun 24, 2025 by thisiscam Loading…
[Pallas][Mosaic GPU] Expose tcgen05.commit_arrive as a standalone primitive.
#29673 opened Jun 23, 2025 by copybara-service bot Loading…
@justinjfu
Add Hermetic C++ Toolchains for JAX project.
#29672 opened Jun 23, 2025 by copybara-service bot Loading…
Expose CompiledMemoryStats::peak_memory_in_bytes to Python.
#29668 opened Jun 23, 2025 by copybara-service bot Loading…
[JAX] Use experimental DCN transfer library as a fallback for PjRt-IFRT cross-host device transfers when the PjRt plugin doesn't implement the cross-host transfer APIs.
#29667 opened Jun 23, 2025 by copybara-service bot Loading…
@emilyfertig
Bump hypothesis from 6.102.4 to 6.135.14 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#29666 opened Jun 23, 2025 by dependabot bot Loading…
@jakevdp
[Mosaic] Add tpu.enqueue_indirect_dma and verification tests
#29665 opened Jun 23, 2025 by copybara-service bot Loading…
Bump fonttools from 4.51.0 to 4.58.4 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#29663 opened Jun 23, 2025 by dependabot bot Loading…
Add a test for bf16 1D memref load -> broadcast issue.
#29661 opened Jun 23, 2025 by copybara-service bot Loading…
[CI] Add bazel TPU presubmit testing - WIP CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#29660 opened Jun 23, 2025 by MichaelHudgins • Draft
Add a method to xla::PjRtClient to indicate whether cross-host device transfers are supported.
#29659 opened Jun 23, 2025 by copybara-service bot Loading…
@emilyfertig
[Mosaic GPU] Add preliminary support for scaled tcgen05.mma
#29656 opened Jun 23, 2025 by copybara-service bot Loading…
@apaszke
Add use_spmd_partitioner in tf.tpu.XlaOptions, TPUReplicateMetadata. The default value is false.
#29646 opened Jun 23, 2025 by copybara-service bot Loading…
[Mosaic GPU] Enable warpgroup semantics for inline_mgpu
#29645 opened Jun 23, 2025 by copybara-service bot Loading…
Move Mosaic into XLA
#29643 opened Jun 23, 2025 by copybara-service bot Loading…
[Mosaic] Fold load->reshape and reshape->store into strided load/store respectively.
#29640 opened Jun 22, 2025 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.