ROCm/jax

Pull Requests Commits

test: strip wheel_tests_continuous.yml to TheRock-only for manual testing

mminutoli committed 1 day ago

b8b504e5

[ROCm] Remove runtime rocm-sdk init; paths now baked in container

mminutoli committed 1 day ago

4f5f31dd

[ROCm] Run rocm-sdk init in ROCm Info step when available

mminutoli committed 1 day ago

44d4bfdf

[ROCm] Unify TheRock into ROCm matrix instead of separate jobs

mminutoli committed 1 day ago

44684571

[ROCm] Use TheRock manylinux container for artifact builds

mminutoli committed 1 day ago

3723d20e

[ROCm] Add TheRock CI integration to nightly and continuous workflows

mminutoli committed 1 day ago

60230965

[ROCm] fix: wire up clone_main_xla for ROCm builds and tests

mminutoli committed 1 day ago

03b39742

Merge pull request #36346 from mattjj:remove-split-transpose

Google-ML-Automation committed 5 days ago

eb95074f

Set xla_tpu_enable_mlir_verifier in tests.

megrez-yliu committed 5 days ago

1bf3fa8e

[fixit] remove vestigial `split_transpose` parameter from `scan`

mattjj committed 5 days ago

01d48465

[XLA:TPU] Abort early if the peak HLO memory consumption is greater than HBM

jcai19 committed 5 days ago

917f27f7

Replace uses of `.vma` (ShapedArray property) in favor of `.manual_axis_type.varying`.

rdyro committed 5 days ago

b03ab674

Merge pull request #36353 from jax-ml:shmap_mat

Google-ML-Automation committed 5 days ago

1ad722d3

Replace vma with mat

yashk2810 committed 5 days ago

8b8ba83e

Expose all_gather_invariant via `jax.lax.all_gather(..., to='invarying')` so that users have a way to go to invarying when the outside shard_map mesh context is Auto since 'reduced' doesn't work in Auto mode.

yashk2810 committed 5 days ago

c99d000a

Merge pull request #36344 from mattjj:remove-scan-linear

Google-ML-Automation committed 5 days ago

fdc12653

[fixit] remove vestigial `linear` param from `scan`

mattjj committed 5 days ago

fdb48fa0

Raise a better error if new_ref/empty_ref is used in eager shard_map. Fixes https://github.com/jax-ml/jax/issues/36000

yashk2810 committed 5 days ago

c7a0eb5e

roll-back #36265

Jake VanderPlas committed 5 days ago

8b2d2311

Fix unmentioned and valid_repeats to include unreduced/reduced too.

yashk2810 committed 5 days ago

e0009378

Propagate sharding and manual type for out type.

Kevin Wang committed 5 days ago

a0a928f3

Remove 1-sized unreduced/reduced mesh axis if `remove_size_one_mesh_axis_from_type` config is True.

yashk2810 committed 5 days ago

44bb9613

Finalize deprecation of `jax.interpreters.ad.reducing_transposes` for JAX v0.10.0.

Jake VanderPlas committed 5 days ago

fdc884ad

Merge pull request #36330 from jakevdp:mlir-dep

Google-ML-Automation committed 5 days ago

813d4bcc

Finalize deprecation of several jax.core APIs for JAX v0.10.0.

Jake VanderPlas committed 5 days ago

db7dce04

Simplify SemanticallyEqualShardings implementation

yashk2810 committed 5 days ago

f59e860a

clean up symbols from mlir.custom_call deprecation

jakevdp committed 5 days ago

4eed0812

Merge pull request #36320 from hawkinsp:uv

Google-ML-Automation committed 5 days ago

4fa6ce78

[Mosaic GPU] Add `tcgen05_commit_arrive` to the MGPU dialect.

allanrenucci committed 5 days ago

88f51af9

Expose `TraceTag` and `set_current_trace` in `jax.extend.core`.

Jake VanderPlas committed 5 days ago

375aeed2

Older