ROCm/jax

Pull Requests Commits

Switch to hermetic llvm for rocm

alekstheod committed 4 days ago

f56f634a

Skip newly added gpu_pallas_interpret_distributed_test on TSAN+ASAN due to timeouts.

belitskiy committed 10 days ago

d8e5aa3e

Automated Code Change

Google-ML-Automation committed 10 days ago

8f91f4eb

[Mosaic GPU] Implement `CollapseShape` expression to propagate SMEM transforms through `memref.collapse_shape`.

bchetioui committed 10 days ago

c5d9cb60

Update XLA dependency to use revision http://github.com/openxla/xla/commit/6530bfb20408d49a63166cb06ee4a01930652cd8

Google-ML-Automation committed 10 days ago

89ed5eed

Replace pe.trace_to_jaxpr_dynamic with pe.trace_to_jaxpr in ad_checkpoint.py

Google-ML-Automation committed 11 days ago

59d0f1c0

Reverts 6fe3e4626169a4bb30cb87526599f5d06b7c467a

rdyro committed 11 days ago

001c1947

Re-disable `test_remote_profiler_gcs_path` test

Matt-Hurd committed 11 days ago

5a042588

Merge pull request #37334 from Saumay:add-scipy-linalg-fiedler

Google-ML-Automation committed 11 days ago

a125378b

Add jax.scipy.linalg.fiedler

Saumay committed 11 days ago

888fe244

Replace pe.trace_to_jaxpr_dynamic with pe.trace_to_jaxpr in ad_checkpoint.py

mattjj committed 11 days ago

dcd0eb59

[Pallas] Add pretty printing for device_ids in semaphore_signal.

sharadmv committed 11 days ago

ae0174e4

Fix return type annotation of HiPspec.to_lo

yashk2810 committed 11 days ago

a4afbf78

[pallas-mpmd] Remove `mpmd_meshes` from the lowering context in favor of meshes on avals.

rdyro committed 11 days ago

6fe3e462

Enable JAX async collective tests on TPU.

mwhittaker committed 11 days ago

27579cba

[Mosaic:GPU] Rollback: Use NCCL API on the host to collective peer addresses.

PatriosTheGreat committed 11 days ago

4fe53bc3

Merge pull request #37333 from Saumay:add-scipy-linalg-companion

Google-ML-Automation committed 11 days ago

f2f35728

* Rename `auto_pvary` to `auto_pcast` since it disables more than just pvarys.

yashk2810 committed 11 days ago

561ec494

Fix empty logs.zip error in Upload Workflow Metadata

danielsuo committed 11 days ago

1d2348a4

Fixes for JAX Bzlmod builds

meteorcloudy committed 11 days ago

4dcbf1af

[pallas:sc] Do not use `pl.pallas_call` in tests

superbobry committed 11 days ago

0838ca5a

Skip SparseCore bounds checker on TPU v7x.

belitskiy committed 11 days ago

fad3f818

Fix the recently added hardware_utils statvfs test on Windows.

belitskiy committed 11 days ago

817b7b23

[Mosaic GPU] Add support for direct f8<->bf16 conversion (added in PTX 9.2)

apaszke committed 11 days ago

5b07ce54

[Mosaic TPU] Relax second minor divisibility constraints on Pallas pipeline windows

apaszke committed 11 days ago

13409e89

[Mosaic GPU] Relax the divisibility constraint for loads/stores with gather indices.

golechwierowicz committed 11 days ago

5f6c4667

[Mosaic GPU] Clean up the PTX ISA version conditionals

apaszke committed 11 days ago

c95506b2

Merge pull request #37536 from gnecula:fix_doc

Google-ML-Automation committed 11 days ago

8c510b28

[Pallas] Correct misspecified `num_mxus` for TPU v3 and v6e in `tpu_info.py`.

bchetioui committed 11 days ago

8eb649d8

[docs] Fix doc build

gnecula committed 11 days ago

aa7c1a6b

Older