-
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
Support For Roofline Analysis Of Pallas Kernels
#32281
opened Oct 2, 2025 by
copybara-service
bot
Loading…
updated Oct 2, 2025
[Mosaic] Make matmul more generally usable by supporting non-contiguous non-contracting dims.
#32021
opened Sep 22, 2025 by
copybara-service
bot
Loading…
updated Oct 2, 2025
[JAX] Make colocated Python's serialization pickle each shared object only once
#32271
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 2, 2025
Update
rules_ml_toolchain to newer version with latest updates and SYCL/OneAPI support
#32273
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 2, 2025
Add
remote_bytes_transferred to Pallas CostEstimate dataclass.
#31760
opened Sep 11, 2025 by
copybara-service
bot
Loading…
updated Oct 2, 2025
Convert IFRT's kString dtype to NumPy's variable length string dtype.
#32274
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 2, 2025
Add support for GCS paths in jax.collect_profile
#32269
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
Update
rules_ml_toolchain to a newer version with SYCL/ONEAPI hermetic build support.
#32110
opened Sep 25, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
Speed up unpacking
.tar.xz files in the hermetic toolchain by using multithreaded xz execution.
#32262
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
Expose
get_concrete_mesh to users through jax.sharding.
#32263
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add an SMEM inference rule for
memref.subview.
#32260
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add SMEM inference rule for
mgpu.custom_op.
#32259
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add SMEM inference rules for
wgmma and tcgen05.mma
#32258
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add a trivial SMEM inference rule for
memref.view.
#32257
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add an SMEM inference pass-through rule for
memref.cast.
#32256
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
Test dumping HLO when reading compilation cache
#31328
opened Aug 26, 2025 by
olupton
Loading…
updated Oct 1, 2025
[Mosaic GPU] Properly handle SMEM in
scf.for.
#32255
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Implement SMEM inference rules for
mgpu.async_load/mgpu.async_store.
#32254
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add SMEM inference rules for
vector.load/vector.store.
#32253
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU][NFC] Specify the memory space for two existing tests.
#32250
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add SMEM inference rules for
with_transforms and slice_smem.
#32249
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
Support For Roofline Analysis Of Pallas Kernels
#32245
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Handle SMEM in the
tmem_alloc layout inference rule.
#32239
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
[Mosaic GPU] Add an SMEM inference rule for
memref.Load/memref.Store.
#32238
opened Oct 1, 2025 by
copybara-service
bot
Loading…
updated Oct 1, 2025
Previous Next
ProTip!
Type g p on any issue or pull request to go back to the pull request listing page.