ROCm/jax

Pull Requests Commits

Disable keep_going as the rocm rbe cluster gets overloaded

alekstheod committed 96 days ago

cb0f6d4d

Fix reshape sharding rule to handle merging and splitting in presence of singleton dimensions

yashk2810 committed 97 days ago

41220ca5

[Mosaic GPU] Check collective attribute is consistent across tcgen05 operations.

allanrenucci committed 97 days ago

e4f757d0

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

Google-ML-Automation committed 97 days ago

73f7e9df

[jaxlib] Fixed the type of `TiledLayout.registers_element_type`

superbobry committed 97 days ago

7273582f

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

Google-ML-Automation committed 98 days ago

88f2faec

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

Google-ML-Automation committed 99 days ago

5a780a03

Merge pull request #36094 from froystig:rm-slab

Google-ML-Automation committed 99 days ago

828c25d0

remove jax.experimental.slab

froystig committed 99 days ago

c53a6d80

update function name in test to match test case name

Google-ML-Automation committed 99 days ago

4d9bdd47

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

yashk2810 committed 99 days ago

cb26766b

Add version guard to qr multiply test.

danielsuo committed 99 days ago

297da810

[Mosaic TPU] Support broadcast with 1D tilings.

yueshengys committed 99 days ago

27e8f0a1

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

emilyfertig committed 99 days ago

a231d35f

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

Google-ML-Automation committed 99 days ago

5957917d

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

belitskiy committed 99 days ago

6fc1a5b1

Merge pull request #36082 from gnecula:api_test

Google-ML-Automation committed 99 days ago

dbdd474b

Add GitHub Action to verify pull requests are squashed.

Google-ML-Automation committed 99 days ago

3d0049e4

Expand assertion message to help debug test flakes

gnecula committed 99 days ago

2102bb91

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

superbobry committed 99 days ago

8dee8fdc

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

bchetioui committed 99 days ago

13dc40d0

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

bchetioui committed 100 days ago

521474d7

Reverts f39a8adc0abfcaf6af0039a637f1851d6298ef6f

Google-ML-Automation committed 100 days ago

4438dede

Add avals to TypedNdArray and TypedInt/TypedFloat/TypedComplex.

hawkinsp committed 100 days ago

f39a8adc

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

allanrenucci committed 100 days ago

c028eb9f

Merge pull request #35104 from jpbrodrick89:jpb/ormrqr

Google-ML-Automation committed 100 days ago

23725340

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

allanrenucci committed 100 days ago

9166e6b6

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

superbobry committed 100 days ago

871364ab

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

golechwierowicz committed 100 days ago

9c64c507

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

Google-ML-Automation committed 100 days ago

8b31887d

Older