Skip to content

Pull requests: jax-ml/jax

New pull request New
482 Open 18,708 Closed
482 Open 18,708 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

Get rid of HloModuleProto from CompiledMemoryStats
#28504 opened May 4, 2025 by copybara-service bot Loading…
Major deps udpate:
#28497 opened May 4, 2025 by copybara-service bot Loading…
[Pallas][Mosaic GPU] Expand TMEM support.
#28495 opened May 2, 2025 by copybara-service bot Loading…
@justinjfu
[Pallas] Allow more int casting tests.
#28494 opened May 2, 2025 by copybara-service bot Loading…
[JAX] Clean up JAX_IFRT_VERSION_NUMBER < 2 since we are well past it
#28493 opened May 2, 2025 by copybara-service bot Loading…
[Mosaic][SC] Expose control over the number of active SC cores
#28492 opened May 2, 2025 by copybara-service bot Loading…
When registering plugins, allow to postpone the creation of the options dict.
#28491 opened May 2, 2025 by copybara-service bot Loading…
[Mosaic GPU] Print instead of warning when skipping flash_attention test
#28485 opened May 2, 2025 by andportnoy Loading…
Add workflow for testing JAX against NumPy nightly builds.
#28484 opened May 2, 2025 by copybara-service bot Loading…
@nitins17
2
create a mosaic_gpu_kernels cc_library and include it as part of the //third_party/py/jax/jaxlib/cuda:cuda_gpu_kernels
#28483 opened May 2, 2025 by copybara-service bot Loading…
[Pallas][Mosaic GPU] Refactor carry_coroutine in warp specialized pipeline to be a callback instead of a coroutine.
#28475 opened May 2, 2025 by copybara-service bot Loading…
@justinjfu
Support np.asarray for JAX arrays with fully replicated shardings when some hosts have no local shards. Only hosts with local shards can fetch np arrays.
#28470 opened May 1, 2025 by copybara-service bot Loading…
@emilyfertig
Increase tolerance in mgpu_attention_test .
#28468 opened May 1, 2025 by copybara-service bot Loading…
@jburnim
[Mosaic GPU] Deallocate TMEM using a single warp CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits kokoro:force-run pull ready Ready for copybara import and testing
#28461 opened May 1, 2025 by andportnoy Loading…
Rewrite broadcast_in_dim to a reshard if only the sharding is changing
#28455 opened May 1, 2025 by dfm Loading…
@dfm
5
Rename backend.compile to backend.compile_and_load.
#28451 opened May 1, 2025 by copybara-service bot Loading…
@danielsuo
[Mosaic GPU] Improve checks in utils.ptr_to_memref
#28450 opened May 1, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add support for arbitrary reductions of tiled layouts
#28448 opened May 1, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic] Support squeezing tiled memrefs to 1d shapes.
#28444 opened May 1, 2025 by copybara-service bot Loading…
Add direct pypi dependencies to the JAX test targets.
#28443 opened May 1, 2025 by copybara-service bot Loading…
[JAX] Add python3.14 wheel test configs
#28440 opened Apr 30, 2025 by kanglant • Draft
Add python bindings for profile data class to JAX Profiler.
#28435 opened Apr 30, 2025 by copybara-service bot Loading…
Enable JAX Profiler tests for TPU's
#28432 opened Apr 30, 2025 by copybara-service bot Loading…
Add type annotations to core abstractify functions.
#28423 opened Apr 30, 2025 by copybara-service bot Loading…
@jakevdp
Only run relevant GPU tests in the optional GPU CI CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28421 opened Apr 30, 2025 by copybara-service bot Loading…
@apaszke
ProTip! Updated in the last three days: updated:>2025-05-02.