Skip to content

Pull requests: jax-ml/jax

New pull request New
605 Open 23,062 Closed
605 Open 23,062 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

Add Windows and build artifacts jobs to mandatory presubmit.
#33712 opened Dec 3, 2025 by copybara-service bot Loading…
@ybaturina
1
add ensure_compile_time_eval support to custom_jvp/vjp kokoro:force-run pull ready Ready for copybara import and testing
#33711 opened Dec 3, 2025 by mattjj Loading…
1
@mattjj
2
[no-thunks] Deprecate old form of cond and remove some tracing layers. pull ready Ready for copybara import and testing
#33710 opened Dec 3, 2025 by dougalm Loading…
6
[Mosaic TPU] Support 1d tiling for packed dtypes when transposing major/minor dims.
#33709 opened Dec 3, 2025 by copybara-service bot Loading…
don't warn on complex->real cast in dot transpose kokoro:force-run pull ready Ready for copybara import and testing
#33708 opened Dec 3, 2025 by mattjj Loading…
1
@mattjj
4
Remove lax_metal_test and associated CI job Apple GPU (Metal) plugin
#33705 opened Dec 3, 2025 by jakevdp Loading…
@jakevdp
2
[Pallas] Skip interpret tests on accelerators
#33704 opened Dec 3, 2025 by copybara-service bot Loading…
@justinjfu
[Mosaic GPU][NFC] Clarify comments on custom call variants.
#33703 opened Dec 3, 2025 by copybara-service bot Loading…
Remove unused imports.
#33700 opened Dec 3, 2025 by copybara-service bot Loading…
@jakeharmon8
Add jax_check_static_indices configuration
#33698 opened Dec 3, 2025 by jakevdp Loading…
@jakevdp
4
Fix lint errors pull ready Ready for copybara import and testing
#33697 opened Dec 3, 2025 by dougalm Loading…
2
[pmap] Remove jax_pmap_shmap_merge configuration.
#33695 opened Dec 3, 2025 by copybara-service bot Loading…
@danielsuo
[Mosaic GPU] Add support for all kinds of TMA reductions.
#33692 opened Dec 3, 2025 by copybara-service bot Loading…
Change int2 shape in pallas test
#33690 opened Dec 3, 2025 by copybara-service bot Loading…
[Pallas:SC] Add plsc.sort_key_val to give access to the mask and descending args.
#33683 opened Dec 3, 2025 by copybara-service bot Loading…
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
[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] 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…
ProTip! Exclude everything labeled bug with -label:bug.