ROCm/jax

Pull Requests Commits

Use dynamic spawn strategy for rocm rbe ci test execution

alekstheod committed 5 days ago

3109b801

Merge pull request #36023 from jax-ml:postrelease/0.9.2

Google-ML-Automation committed 5 days ago

756c1bc5

[Mosaic GPU] Erase lowered operations immediately.

allanrenucci committed 5 days ago

114857b1

[pallas:mosaic] Added `CoreMemorySpace`

superbobry committed 5 days ago

2fe54d1c

Merge remote-tracking branch 'origin/release/0.9.2' into postrelease/0.9.2

danielsuo committed 5 days ago

a7d126ee

[Pallas MGPU] Fix the abstract eval of semaphore_signal_multicast to support indexing.

Rifur13 committed 5 days ago

ac27e9ec

[Mosaic GPU] Add `WarpMapOp` to the dialect.

allanrenucci committed 5 days ago

a5717b2f

[pallas:mosaic] Refined the `NotImplementedError` condition in `_device_id_to_logical`

superbobry committed 5 days ago

fe0508a5

Merge pull request #35611 from ROCm:araganes/skip-pmap-rocm

Google-ML-Automation committed 5 days ago

528bc5e0

Merge pull request #26604 from ROCm:min_max_atomics

Google-ML-Automation committed 5 days ago

55117a99

Merge pull request #35917 from ROCm:skip-tridiagonal-solve-grad0-rocm

Google-ML-Automation committed 5 days ago

47459c16

[Mosaic] Define verifyLayout for TiledLayoutAttr

tlongeri committed 5 days ago

5e51d290

[Pallas:MGPU] Create helper function to make an operation conform to the `IsolatedFromAbove` trait.

allanrenucci committed 5 days ago

37107d62

[Pallas] Document the new features added in the `0.9.2` release for Pallas/Mosaic GPU.

bchetioui committed 5 days ago

be297dd9

Merge pull request #36027 from superbobry:pyrefly-2

Google-ML-Automation committed 5 days ago

3a206dda

[pallas] Preserve concrete bounds in `pl.loop`

superbobry committed 5 days ago

7e81f45c

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

Google-ML-Automation committed 5 days ago

31c3a650

[pallas:mosaic] Do not handle JAX memory spaces in the lowering

superbobry committed 5 days ago

468a1e9a

Merge pull request #36012 from jakevdp:pyrefly-057

Google-ML-Automation committed 5 days ago

ec97111d

Allow `with_refs` to take Accums too.

yashk2810 committed 5 days ago

474cc95f

[pallas:mosaic] `pltpu.CoreType` now uses strings as values

superbobry committed 5 days ago

00d408c5

jnp.arange: avoid constant-folding when using step

Jake VanderPlas committed 5 days ago

3dd61872

Make sure that clear_backends() actually clears all the caches.

pschuh committed 5 days ago

b828953b

[ci] update to pyrefly v0.57.0

jakevdp committed 5 days ago

599e8d7f

Add `to_ct_mt` on ManualAxisType to mirror PartitionSpec. Might be helpful for HiJAX (TBD though)

yashk2810 committed 5 days ago

989281dc

Prepare for JAX release 0.9.2

danielsuo committed 5 days ago

a659757d

Cherry-pick upstream XLA and Shardy fixes via Bazel patches.

danielsuo committed 5 days ago

e0bb88af

Preserve existing mhlo.frontend_attributes when adding compute_type annotations.

Google-ML-Automation committed 5 days ago

082f715c

PyClient::CompileAndLoad is already cloning the mlir modules.

pschuh committed 5 days ago

345f336f

Merge pull request #35984 from mattjj:qdd-fix

Google-ML-Automation committed 6 days ago

347982d0

Older