-
Notifications
You must be signed in to change notification settings - Fork 3.3k
Pull requests: jax-ml/jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Change PjRt to use new copy of coordination service.
#33675
opened Dec 2, 2025 by
copybara-service
bot
Loading…
[Pallas] Make rng tests inherit from JaxTestCase
#33672
opened Dec 2, 2025 by
copybara-service
bot
Loading…
[JAX] Split _jax extension into _xla and _jax extensions.
#33671
opened Dec 2, 2025 by
copybara-service
bot
Loading…
Fix breakage cause by list being passed to psum
#33669
opened Dec 2, 2025 by
copybara-service
bot
Loading…
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…
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…
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…
Add RPCs and client support for continuous profiling snapshots
#33644
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…
[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…
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…
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…
Previous Next
ProTip!
no:milestone will show everything without a milestone.