Skip to content

Pull requests: jax-ml/jax

New pull request New
559 Open 22,758 Closed
559 Open 22,758 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

[hijax] add optional JVP rule to VJPHiPrimitive kokoro:force-run pull ready Ready for copybara import and testing
#33307 opened Nov 13, 2025 by mattjj Loading…
@mattjj
6
[Mosaic GPU] Add an assertion to check for remaining constraints after solving.
#33306 opened Nov 13, 2025 by copybara-service bot Loading…
[Mosaic GPU] Use a recursive tree implementation for vector_concat
#33305 opened Nov 13, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Tune the packing length for i4 -> f8e4m3fn conversion
#33304 opened Nov 13, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Use the packing trick for int4 -> bf16 conversions
#33303 opened Nov 13, 2025 by copybara-service bot Loading…
@apaszke
[mosaic] Made GetMosaicHardwareGeneration public
#33301 opened Nov 13, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Support barrier.arrive with orders_tensor_core in WG semantics.
#33298 opened Nov 13, 2025 by copybara-service bot Loading…
Automated Code Change
#33293 opened Nov 13, 2025 by copybara-service bot Loading…
[pmap] Implement device_put_sharded using NamedSharding.
#33292 opened Nov 13, 2025 by copybara-service bot Loading…
@danielsuo
Add CUDA UMD libs to the runfiles of the tests to enable RBE testing of JAX CUDA 13 wheels in RBE.
#33291 opened Nov 13, 2025 by copybara-service bot Loading…
@ybaturina
1
[pmap] Deprecate jax.sharding.PmapSharding, jax.device_put_replicated and jax.device_put_sharded.
#33290 opened Nov 13, 2025 by copybara-service bot Loading…
@danielsuo
Allow setting jax_num_cpu_devices to the same value even after backend initialization.
#33288 opened Nov 12, 2025 by copybara-service bot Loading…
1
@hawkinsp
[JAX] Re-enable grad(sin(x)) test case in xla_metadata_test.
#33287 opened Nov 12, 2025 by copybara-service bot Loading…
Align/correlate Mosaic inject xplane to cupti kernel execution event.
#33286 opened Nov 12, 2025 by copybara-service bot Loading…
Add TraceMe for Mosaic Kernel Injection Id
#33285 opened Nov 12, 2025 by copybara-service bot Loading…
[Pallas] Implement a "rolled scheduler" for the pipelining API.
#33272 opened Nov 12, 2025 by copybara-service bot Loading…
@justinjfu
Update pull/push rule API for call_hi_primitive_p
#33266 opened Nov 12, 2025 by copybara-service bot Loading…
[Pallas/Mosaic TPU] Allow non-leading and non-matching batch dimensions in dot_general.
#33265 opened Nov 11, 2025 by copybara-service bot Loading…
Derive the number of SC cores and subcores from hardware info by default.
#33263 opened Nov 11, 2025 by copybara-service bot Loading…
@IvyZX
[Mosaic] Allow padding in small tiling row shuffle reshape.
#33262 opened Nov 11, 2025 by copybara-service bot Loading…
[Mosaic] Internal change.
#33261 opened Nov 11, 2025 by copybara-service bot Loading…
fix: typo training book docs for adam-update and init-adam-state
#33254 opened Nov 11, 2025 by pevers Loading…
2
Add import tests for JAX wheels.
#33246 opened Nov 11, 2025 by copybara-service bot Loading…
Expose client's platform version via jax.backend_version(platform)
#33238 opened Nov 11, 2025 by gspschmid Loading…
4
Aot
#33225 opened Nov 10, 2025 by danielsuo Loading…
7
ProTip! Type g i on any issue or pull request to go back to the issue listing page.