ROCm/jax

Pull Requests Commits

Upgrade upstream ROCm CI from 7.2.0 to 7.2.2

Ruturaj4 committed 7 days ago

15ed381b

[Pallas] Add .like constructor method for MemoryRefs in Pallas TPU and GPU

sharadmv committed 7 days ago

d24ffbd0

Refactor Pallas MPMD tests into a dedicated file.

sharadmv committed 7 days ago

9c291a71

Add a single-chip case for TPU 8i to mesh_utils.py.

hawkinsp committed 7 days ago

495c0664

Add PRNG key support to colocated_python.

ZacharyGarrett committed 8 days ago

a4ab391c

[Pallas/Mosaic GPU] Remove `TransposeTransform` from public API and clean up related lowerings.

bchetioui committed 8 days ago

923aa780

CUTE DSL notebooks cleanup and tolerance relaxation.

rdyro committed 8 days ago

f57860db

[cleanup] Remove TPU ragged dot python lowering transpose which is no longer needed, since the C++ path handles it.

rdyro committed 8 days ago

70b4c3d6

Add a name to the Bazel dependency violations job.

belitskiy committed 8 days ago

e6289a82

[Take 2] Move PartitionSpec back to python and intern it.

yashk2810 committed 8 days ago

5e2e7fba

Merge pull request #37237 from samanklesaria:ball_out_sharding

Google-ML-Automation committed 8 days ago

65a8aff8

[Pallas MGPU] Add new MGPU semaphore primitives that take a parameter that specifies the memory scope.

Rifur13 committed 8 days ago

b21b7e22

[Pallas] Enable mpmd with pl.kernel

sharadmv committed 8 days ago

108f08d2

Don't call clean_up_dead_vars in mlir.py.

hawkinsp committed 8 days ago

2f9c65ae

Merge pull request #37072 from ROCm:switch_to_latest_rules_ml_toolchain

Google-ML-Automation committed 8 days ago

1808cfa0

Remove deleted targets from the Bazel dependency Violations workflow.

belitskiy committed 8 days ago

ce93599c

Some tiny optimizations to TypedInt, TypedFloat and TypedComplex:

hawkinsp committed 8 days ago

849fa5b7

Remove _manual_axis_type from `empty` now that it's unused given pl.kernel moved to mpmd_map

yashk2810 committed 8 days ago

e11e9d2e

Move PartitionSpec back to python and intern it.

Google-ML-Automation committed 8 days ago

7da762c9

[CI] Add //jaxlib/mosaic:tiled_layout_test to exceptions in Bazel graph workflow.

belitskiy committed 8 days ago

fc51a014

propagate debug_info arg_names from hi-jaxprs to lo-jaxprs

Cristian Garcia committed 8 days ago

f4455bbb

Update rules_ml_toolchain reference

alekstheod committed 8 days ago

aaf48382

Move PartitionSpec back to python and intern it.

hawkinsp committed 8 days ago

72b90f40

[Pallas/Mosaic GPU] Implement logic that allows commuting `ReshapeTransform`s with `UntilingTransform`s.

bchetioui committed 8 days ago

4cda3a9d

[Mosaic GPU] Require explicit indexing for barriers initialised with num_barriers > 1.

allanrenucci committed 8 days ago

15371b08

[random] Add out_sharding to jax.random.ball

samanklesaria committed 8 days ago

88d679c3

[Mosaic GPU] Disable XLA autotuning for `matmul_test`.

bchetioui committed 8 days ago

3345b109

[pallas:sc] Added `plsc.Indices` which allow specifying the ignored offset value for gathers/scatters

superbobry committed 8 days ago

a7ac409f

[Mosaic GPU] Allow collapsing unit dimensions when issuing TMA loads.

bchetioui committed 8 days ago

a56a2210

[CI] Fix portserver not being shut down on the doc render Sphinx command failing.

belitskiy committed 8 days ago

59e12d5e

Older