Skip to content

Pull requests: jax-ml/jax

New pull request New
603 Open 23,038 Closed
603 Open 23,038 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

[no-thunks] Reduce stack frames in cond and friends
#33677 opened Dec 2, 2025 by dougalm Loading…
6
Change PjRt to use new copy of coordination service.
#33675 opened Dec 2, 2025 by copybara-service bot Loading…
@mwhittaker
[Pallas] Make rng tests inherit from JaxTestCase
#33672 opened Dec 2, 2025 by copybara-service bot Loading…
@justinjfu
[JAX] Split _jax extension into _xla and _jax extensions.
#33671 opened Dec 2, 2025 by copybara-service bot Loading…
@hawkinsp
Fix breakage cause by list being passed to psum
#33669 opened Dec 2, 2025 by copybara-service bot Loading…
@yashk2810
add example usage to log_compiles
#33665 opened Dec 2, 2025 by cgarciae Loading…
4
Fix the jax line search, the zoom may only store two points in lo, hi and rec.
#33664 opened Dec 2, 2025 by copybara-service bot Loading…
[Pallas:MGPU] Add support for jnp.sin
#33662 opened Dec 2, 2025 by copybara-service bot Loading…
@apaszke
1
Add constructor to NamedSharding accepting axis names.
#33659 opened Dec 2, 2025 by copybara-service bot Loading…
(2/N) Add support for NamedSharding in existing HloShardingUtil methods. Remaining methods will be updated in follow up cl's.
#33658 opened Dec 2, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add support for reductions when writing to GMEM via TMA in WG semantics.
#33655 opened Dec 2, 2025 by copybara-service bot Loading…
[Mosaic GPU] Optimize the computation of tcgen05.mma matrix descriptors
#33653 opened Dec 2, 2025 by copybara-service bot Loading…
@apaszke
This is a mechanical change to replace tile_assignment().dimensions() by dimensions().
#33651 opened Dec 2, 2025 by copybara-service bot Loading…
This is a mechanical change to replace tile_assignment().num_dimensions() by num_dimensions().
#33650 opened Dec 2, 2025 by copybara-service bot Loading…
Integrate Triton up to [8d445186](https://github.com/openai/triton/commits/ce7393e1e407a84513f6f133f93791878d445186)
#33647 opened Dec 2, 2025 by copybara-service bot Loading…
[Pallas MGPU] Use the cp.async.bulk instruction for large contiguous copies. Currently we’re limited to 256 elements per dimension when using the tensormap in cp.async.bulk.tensor.
#33645 opened Dec 2, 2025 by copybara-service bot Loading…
@Rifur13
Add RPCs and client support for continuous profiling snapshots
#33644 opened Dec 2, 2025 by copybara-service bot Loading…
Automated Code Change
#33641 opened Dec 2, 2025 by copybara-service bot Loading…
Add test case to debug_print scalar in Pallas.
#33639 opened Dec 2, 2025 by copybara-service bot Loading…
@IvyZX
Support Hijax types in emit_pipeline.
#33638 opened Dec 2, 2025 by copybara-service bot Loading…
[Pallas] Add missing docstrings and clean up mis-renderings. kokoro:force-run pull ready Ready for copybara import and testing
#33632 opened Dec 1, 2025 by justinjfu Loading…
5
Delete the floating point optimization level and memory fitting level settings in favor of the EffortLevel enum.
#33631 opened Dec 1, 2025 by copybara-service bot Loading…
[Pallas] Fix cost estimate einsum bug.
#33628 opened Dec 1, 2025 by copybara-service bot Loading…
@justinjfu
Add CI to run ty on jax/_src/dtypes.py
#33626 opened Dec 1, 2025 by jakevdp Loading…
@jakevdp
4
Replace jax.lax.pvary with jax.lax.pcast(..., to='varying'). No function changes intended. This is just a rename.
#33621 opened Dec 1, 2025 by copybara-service bot Loading…
@yashk2810
ProTip! no:milestone will show everything without a milestone.