ROCm/jax

Pull Requests Commits

update to test if there are GPU events when doing profiling on ROCm

cj401-amd committed 19 days ago

85f42c1e

update function name in test to match test case name

Google-ML-Automation committed 19 days ago

4d9bdd47

Make set_mesh(None) work as a way to unset mesh by setting an empty_concrete_mesh internally

yashk2810 committed 19 days ago

cb26766b

Add version guard to qr multiply test.

danielsuo committed 19 days ago

297da810

[Mosaic TPU] Support broadcast with 1D tilings.

yueshengys committed 19 days ago

27e8f0a1

[Shardy MPMD] Add support for unreduced axes to NamedSpmdShardingSpec.

emilyfertig committed 19 days ago

a231d35f

Merge pull request #36056 from ROCm:rocm-wheel-version-suffix

Google-ML-Automation committed 19 days ago

5957917d

Decrease the number of test cases for sparse_bcoo_bcsr_test on H100 (internal).

belitskiy committed 19 days ago

6fc1a5b1

Merge pull request #36082 from gnecula:api_test

Google-ML-Automation committed 19 days ago

dbdd474b

Add GitHub Action to verify pull requests are squashed.

Google-ML-Automation committed 19 days ago

3d0049e4

Expand assertion message to help debug test flakes

gnecula committed 19 days ago

2102bb91

[pallas:mosaic] It is no longer necessary to use "main" as the name of the kernel

superbobry committed 19 days ago

8dee8fdc

[Pallas/Mosaic GPU][NFC] Add missing type argument to (never executed) call to `mgpu.c`.

bchetioui committed 19 days ago

13dc40d0

[Pallas/Mosaic GPU] Remove pre-0.9.2 version guards.

bchetioui committed 19 days ago

521474d7

Reverts f39a8adc0abfcaf6af0039a637f1851d6298ef6f

Google-ML-Automation committed 19 days ago

4438dede

Add avals to TypedNdArray and TypedInt/TypedFloat/TypedComplex.

hawkinsp committed 19 days ago

f39a8adc

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

allanrenucci committed 19 days ago

c028eb9f

Merge pull request #35104 from jpbrodrick89:jpb/ormrqr

Google-ML-Automation committed 19 days ago

23725340

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

allanrenucci committed 19 days ago

9166e6b6

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

superbobry committed 19 days ago

871364ab

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

golechwierowicz committed 19 days ago

9c64c507

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

Google-ML-Automation committed 19 days ago

8b31887d

[Mosaic GPU] Add AsyncStoreScalesSmemToTmemOp to the dialect.

golechwierowicz committed 19 days ago

0e1d52ce

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

yashk2810 committed 19 days ago

83ac59f1

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

Google-ML-Automation committed 19 days ago

6587096a

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

golechwierowicz committed 19 days ago

28880dc0

Reverts 9cecb18d3d811b751aebfc6c6e8f97caf169be6d

hawkinsp committed 19 days ago

e49045a4

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

golechwierowicz committed 19 days ago

91a687b2

Reverts d1eb0fb0717c58b3bf10e668eb5ed2b4702b2e72

hawkinsp committed 19 days ago

04c2d914

Reverts 6bc1cdb7c8b62df90af526476ac59e2bdb803c1b

hawkinsp committed 19 days ago

238f6c54

Older