ROCm/jax

Pull Requests Commits

add offload-compress flag to reduce rocm wheel sizze

gulsumgudukbay committed 7 hours ago

d73557de

Merge pull request #36044 from ROCm:enable_dynamic_spawn_strategy_for_rocm_rbe

Google-ML-Automation committed 9 hours ago

b23653d9

[Mosaic GPU] Lift restriction on tcgen05.mma sparse metadata WG lowering.

golechwierowicz committed 9 hours ago

e3ee2433

[Mosaic GPU] Throw ValueError for unsupported data types in async_copy_scales_smem_to_tmem.

golechwierowicz committed 9 hours ago

f5cc0e2f

Fix a possible crash in weakref_lru_cache.

hawkinsp committed 9 hours ago

62428e4a

[pmap] Remove `cpp_pmap` from `api.py`.

danielsuo committed 10 hours ago

df68a7ab

[pmap] Remove the config state `jax_pmap_shmap_merge`.

danielsuo committed 10 hours ago

4bc95df7

Use dynamic spawn strategy for rocm rbe ci test execution

alekstheod committed 10 hours ago

3109b801

Disable asan, msan, and tsan tests for tpu_pallas_pipeline_test to avoid timeouts.

danielsuo committed 10 hours ago

a44a1d84

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

Google-ML-Automation committed 10 hours ago

756c1bc5

[Mosaic GPU] Erase lowered operations immediately.

allanrenucci committed 10 hours ago

114857b1

[pallas:mosaic] Added `CoreMemorySpace`

superbobry committed 12 hours ago

2fe54d1c

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

danielsuo committed 12 hours ago

a7d126ee

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

Rifur13 committed 13 hours ago

ac27e9ec

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

allanrenucci committed 14 hours ago

a5717b2f

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

superbobry committed 14 hours ago

fe0508a5

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

Google-ML-Automation committed 14 hours ago

528bc5e0

Merge pull request #26604 from ROCm:min_max_atomics

Google-ML-Automation committed 14 hours ago

55117a99

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

Google-ML-Automation committed 14 hours ago

47459c16

[Mosaic] Define verifyLayout for TiledLayoutAttr

tlongeri committed 15 hours ago

5e51d290

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

allanrenucci committed 15 hours ago

37107d62

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

bchetioui committed 16 hours ago

be297dd9

Merge pull request #36027 from superbobry:pyrefly-2

Google-ML-Automation committed 17 hours ago

3a206dda

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

superbobry committed 17 hours ago

7e81f45c

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

Google-ML-Automation committed 18 hours ago

31c3a650

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

superbobry committed 1 day ago

468a1e9a

Merge pull request #36012 from jakevdp:pyrefly-057

Google-ML-Automation committed 1 day ago

ec97111d

Allow `with_refs` to take Accums too.

yashk2810 committed 1 day ago

474cc95f

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

superbobry committed 1 day ago

00d408c5

jnp.arange: avoid constant-folding when using step

Jake VanderPlas committed 1 day ago

3dd61872

Older