ROCm/jax

Pull Requests Commits

update for slotdet when n=1 or 2

cj401-amd committed 234 days ago

812f2680

[Pallas:MGPU] Fix `plgpu.inline_mgpu` support for scalar inputs under WG semantics.

allanrenucci committed 234 days ago

c028eb9f

Merge pull request #35104 from jpbrodrick89:jpb/ormrqr

Google-ML-Automation committed 234 days ago

23725340

[Pallas:MGPU][NFC] Fix type annotations.

allanrenucci committed 234 days ago

9166e6b6

[pallas:mosaic] Allow passing `use_tc_tiling_on_sc` when using `mpmd_map`

superbobry committed 234 days ago

871364ab

[Mosaic GPU] Add layout inference rule to AsyncStoreScalesSmemToTmem op.

golechwierowicz committed 234 days ago

9c64c507

[Mosaic TPU] Allow non-DMA semaphores in EnqueueDMA for src semaphore as well

Google-ML-Automation committed 234 days ago

8b31887d

[Mosaic GPU] Add AsyncStoreScalesSmemToTmemOp to the dialect.

golechwierowicz committed 234 days ago

0e1d52ce

Do a few more cleanups after the pmap shmap merge deletion.

yashk2810 committed 234 days ago

83ac59f1

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

Google-ML-Automation committed 234 days ago

6587096a

[Mosaic GPU] Add pallas level test for collective sparse tcgen05.mma.

golechwierowicz committed 234 days ago

28880dc0

Reverts 9cecb18d3d811b751aebfc6c6e8f97caf169be6d

hawkinsp committed 234 days ago

e49045a4

[Mosaic GPU] Add dtype to `helpers.format_tcgen05_sparse_metadata`.

golechwierowicz committed 234 days ago

91a687b2

Reverts d1eb0fb0717c58b3bf10e668eb5ed2b4702b2e72

hawkinsp committed 234 days ago

04c2d914

Reverts 6bc1cdb7c8b62df90af526476ac59e2bdb803c1b

hawkinsp committed 234 days ago

238f6c54

Reverts 3c2ae8ea914bb0b49d200fcc00f294c99ebb91e8

hawkinsp committed 234 days ago

af17b87d

Rename `f` to `fun` in `jax.pmap` as it was before.

yashk2810 committed 234 days ago

3c2ae8ea

Move the pmap API from api.py to pmap.py to avoid local imports

yashk2810 committed 234 days ago

6bc1cdb7

Do a few more cleanups after pmap_shmap merge

yashk2810 committed 234 days ago

d1eb0fb0

make new_ref kind= kwarg public, give anselm_ref option a meaningful name.

levskaya committed 234 days ago

9e2e4bdc

[pmap] Remove pmap from jaxlib.

danielsuo committed 234 days ago

fa21532f

Add `RemoveSubAxesInInputOutputShardingsPass` to the export pipeline.

ZixuanJiang committed 234 days ago

7109fea0

Remove xla_pmap_p since it's deleted from JAX

yashk2810 committed 234 days ago

fbea98f8

Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone

yashk2810 committed 234 days ago

9cecb18d

Merge pull request #36046 from jax-ml:lockfile/0.9.2

Google-ML-Automation committed 234 days ago

9d906b0f

Add more async start/done ops to JAX.

mwhittaker committed 234 days ago

3d6b5b79

[pmap] Remove remaining `pmap` references from `jax`.

danielsuo committed 234 days ago

5d3868da

Add ormqr primitive and jax.scipy.linalg.qr_multiply

jpbrodrick89 committed 235 days ago

4e87ba16

Update force forward compat logic.

Google-ML-Automation committed 235 days ago

0679014e

[pallas:mosaic] Wired in `CoreMemorySpace` in the lowering

superbobry committed 235 days ago

aa9977f0

Older