ROCm/jax

Pull Requests Commits

Remove old ROCm build code

charleshofer committed 100 days ago

df4f394d

[pallas:mosaic] Added `CoreMemorySpace`

superbobry committed 100 days ago

2fe54d1c

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

Rifur13 committed 100 days ago

ac27e9ec

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

allanrenucci committed 100 days ago

a5717b2f

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

superbobry committed 100 days ago

fe0508a5

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

Google-ML-Automation committed 100 days ago

528bc5e0

Merge pull request #26604 from ROCm:min_max_atomics

Google-ML-Automation committed 100 days ago

55117a99

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

Google-ML-Automation committed 100 days ago

47459c16

[Mosaic] Define verifyLayout for TiledLayoutAttr

tlongeri committed 100 days ago

5e51d290

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

allanrenucci committed 100 days ago

37107d62

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

bchetioui committed 100 days ago

be297dd9

Merge pull request #36027 from superbobry:pyrefly-2

Google-ML-Automation committed 100 days ago

3a206dda

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

superbobry committed 100 days ago

7e81f45c

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

Google-ML-Automation committed 100 days ago

31c3a650

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

superbobry committed 100 days ago

468a1e9a

Merge pull request #36012 from jakevdp:pyrefly-057

Google-ML-Automation committed 100 days ago

ec97111d

Allow `with_refs` to take Accums too.

yashk2810 committed 100 days ago

474cc95f

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

superbobry committed 100 days ago

00d408c5

jnp.arange: avoid constant-folding when using step

Jake VanderPlas committed 100 days ago

3dd61872

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

pschuh committed 100 days ago

b828953b

[ci] update to pyrefly v0.57.0

jakevdp committed 100 days ago

599e8d7f

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

yashk2810 committed 100 days ago

989281dc

Preserve existing mhlo.frontend_attributes when adding compute_type annotations.

Google-ML-Automation committed 100 days ago

082f715c

PyClient::CompileAndLoad is already cloning the mlir modules.

pschuh committed 100 days ago

345f336f

Merge pull request #35984 from mattjj:qdd-fix

Google-ML-Automation committed 101 days ago

347982d0

[mosaic] Fixed relative references in FAILUREOR_ASSIGN_OR_RETURN

superbobry committed 101 days ago

090e2cae

[Pallas][GPU kernel interpreter] Support interpreting of `copy_gmem_to_smem`.

Google-ML-Automation committed 101 days ago

e637fcb4

[Mosaic GPU] Add TMEM aliasing support for WG semantics.

dimitar-asenov committed 101 days ago

521f7951

Add regression test of https://github.com/jax-ml/jax/pull/35787

GleasonK committed 101 days ago

a231919a

Replace `SDS(vma: frozenset = ...)` to `SDS(manual_type: jax.sharding.ManualAxisType = ...)`. Also expose `jax.sharding.ManualAxisType` as a public API.

yashk2810 committed 101 days ago

2ed3944c

Older