Skip to content

Pull requests: jax-ml/jax

New pull request New
620 Open 23,148 Closed
620 Open 23,148 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

[maint] clean up jnp.arange implementation pull ready Ready for copybara import and testing
#33819 opened Dec 9, 2025 by jakevdp Loading…
@jakevdp
3
[Pallas] Device Id dict to mesh fastpath for power of twos
#33817 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…
[Mosaic GPU] Remove {Least,Most}ReplicatedExpression constructors.
#33809 opened Dec 8, 2025 by copybara-service bot Loading…
@bchetioui
[Mosaic GPU] Support FragmentedArray.broadcast_in_dim for splat to other layouts.
#33808 opened Dec 8, 2025 by copybara-service bot Loading…
Use ifrt::AttributeMap::Get instead of directly accessing map
#33804 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
Bump actions/checkout from 6.0.0 to 6.0.1 dependencies Pull requests that update a dependency file github_actions Pull requests that update Github_actions code
#33802 opened Dec 8, 2025 by dependabot bot Loading…
[Pallas:SC] Adds a test for lowering a custom primitive around a kernel. Modifies pl.kernel to support it properly, using state discharge.
#33800 opened Dec 8, 2025 by copybara-service bot Loading…
Integrate Triton up to 8d445186
#33799 opened Dec 8, 2025 by copybara-service bot Loading…
Fix links in contributing.md to use HTTPS
#33795 opened Dec 8, 2025 by partev Loading…
7
docs: Add tutorial on Attention Mechanism in JAX
#33791 opened Dec 8, 2025 by vibeswithkk Loading…
5 tasks done
9
Automated Code Change
#33789 opened Dec 7, 2025 by copybara-service bot Loading…
Allow overriding the XLA commit to use in Continuous workflow
#33784 opened Dec 6, 2025 by copybara-service bot Loading…
@nitins17
Export pyproject.toml in the BUILD file
#33783 opened Dec 5, 2025 by copybara-service bot Loading…
@superbobry
[jax.collect_profile] Allow arbitrary options to be passed to XProf
#33782 opened Dec 5, 2025 by copybara-service bot Loading…
[Pallas TPU] Allow closed over scalars in core_map code
#33781 opened Dec 5, 2025 by copybara-service bot Loading…
@sharadmv
Add collective to async_copy_scales_to_tmem.
#33780 opened Dec 5, 2025 by copybara-service bot Loading…
[pallas:mosaic] pltpu.emit_pipeline now accepts block specs in HBM
#33776 opened Dec 5, 2025 by copybara-service bot Loading…
@superbobry
Move rules_ml_toolchain up due to request move default rules_ml_toolchain initialization to @xla//:workspace4.bzl
#33774 opened Dec 5, 2025 by copybara-service bot Loading…
1
Colocated python perf optimization.
#33772 opened Dec 5, 2025 by copybara-service bot Loading…
Add initial blackwell mxfp4 blockscale matmul implementation.
#33770 opened Dec 5, 2025 by copybara-service bot Loading…
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.