-
Notifications
You must be signed in to change notification settings - Fork 3.2k
Pull requests: jax-ml/jax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
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…
[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…
[Mosaic GPU] Tune the packing length for i4 -> f8e4m3fn conversion
#33304
opened Nov 13, 2025 by
copybara-service
bot
Loading…
[Mosaic GPU] Use the packing trick for int4 -> bf16 conversions
#33303
opened Nov 13, 2025 by
copybara-service
bot
Loading…
[mosaic] Made
GetMosaicHardwareGeneration public
#33301
opened Nov 13, 2025 by
copybara-service
bot
Loading…
[Mosaic GPU] Support
barrier.arrive with orders_tensor_core in WG semantics.
#33298
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…
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…
[pmap] Deprecate
jax.sharding.PmapSharding, jax.device_put_replicated and jax.device_put_sharded.
#33290
opened Nov 13, 2025 by
copybara-service
bot
Loading…
Allow setting jax_num_cpu_devices to the same value even after backend initialization.
#33288
opened Nov 12, 2025 by
copybara-service
bot
Loading…
[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…
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…
[Mosaic] Allow padding in small tiling row shuffle reshape.
#33262
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…
Expose client's platform version via jax.backend_version(platform)
#33238
opened Nov 11, 2025 by
gspschmid
Loading…
Previous Next
ProTip!
Type g i on any issue or pull request to go back to the issue listing page.