Skip to content

Pull requests: jax-ml/jax

New pull request New
626 Open 23,175 Closed
626 Open 23,175 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

[export] Add backwards compatibility tests for v6. pull ready Ready for copybara import and testing
#33853 opened Dec 10, 2025 by gnecula Loading…
@gnecula
1
[Pallas:MGPU] Properly restore the pytree structure when unflattening unions
#33851 opened Dec 10, 2025 by copybara-service bot Loading…
@apaszke
[Pallas:MGPU] Expose fragmented_array.Replicated as part of the public API.
#33849 opened Dec 10, 2025 by copybara-service bot Loading…
[export] Add backwards compatibility test for memory_space pull ready Ready for copybara import and testing
#33848 opened Dec 10, 2025 by gnecula Loading…
@gnecula
1
jnp.arange: avoid device transfer when possible
#33846 opened Dec 10, 2025 by jakevdp Loading…
@jakevdp
1
Internal change to JAX lib visibility.
#33845 opened Dec 10, 2025 by copybara-service bot Loading…
[pmap] Add default_pmap_sharding to migrate users away from PmapSharding.default.
#33844 opened Dec 10, 2025 by copybara-service bot Loading…
@danielsuo
Automated Code Change
#33843 opened Dec 10, 2025 by copybara-service bot Loading…
Remove of disassemble_into_single_device_arrays in favor of a new
#33840 opened Dec 9, 2025 by copybara-service bot Loading…
@pschuh
Remove linear util use in pjit
#33839 opened Dec 9, 2025 by dougalm Loading…
1
Load nvjitlink library version 12
#33838 opened Dec 9, 2025 by yliu120 Loading…
3
Remove another instance of access to AttributeMap::Map
#33835 opened Dec 9, 2025 by copybara-service bot Loading…
Test TPU v7
#33833 opened Dec 9, 2025 by quoctruong • Draft
1
Make xla::ifrt::AttributeMap thread-safe
#33832 opened Dec 9, 2025 by copybara-service bot Loading…
Aggregate cross-host data transfers made by jax.device_put
#33828 opened Dec 9, 2025 by rao-ashish Loading…
7
[CI] Make build event link easier to find CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#33827 opened Dec 9, 2025 by MichaelHudgins • Draft
4
Test nightly TPU.
#33826 opened Dec 9, 2025 by copybara-service bot Loading…
@ybaturina
1
changes in jax/_src/numpy/reductions.py
#33823 opened Dec 9, 2025 by haarit19058 Loading…
6
[MGPU] Add support for broadcast on major dim in WGStridedFragLayout.
#33822 opened Dec 9, 2025 by copybara-service bot Loading…
Integrate Triton up to 8d445186
#33821 opened Dec 9, 2025 by copybara-service bot Loading…
Expose profiler advanced configuration as a Python dict.
#33816 opened Dec 9, 2025 by copybara-service bot Loading…
Support Hijax types in emit_pipeline.
#33813 opened Dec 8, 2025 by copybara-service bot Loading…
@davisyoshida
Update PyArray::__init__ to also take a PyArray.
#33812 opened Dec 8, 2025 by copybara-service bot Loading…
@pschuh
Reverts 32d830f9cf213f37a50a058cabda2ab8b230ee7f
#33811 opened Dec 8, 2025 by copybara-service bot Loading…
[pmap] Update uses of PmapSharding.default with default_pmap_sharding because PmapSharding is deprecated.
#33803 opened Dec 8, 2025 by copybara-service bot Loading…
@danielsuo
ProTip! Type g i on any issue or pull request to go back to the issue listing page.