ROCm/jax

Pull Requests Commits

[0.6.0-UT] Skip test_dot for ROCm <6.5, and shapes exceed shared memory limits

zahiqbal committed 151 days ago

42c611e5

Select the platform GPUs for a few MultiProcessGpuTest test cases.

Pham Binh committed 151 days ago

5822d672

Remove warnings.warn from get_rocm_version

psanal35 committed 151 days ago

48e290b6

Update tests/lax_numpy_test.py

gulsumgudukbay committed 151 days ago

dec5a05a

Update tests/experimental_rnn_test.py

gulsumgudukbay committed 151 days ago

dda0aa17

[0.6.0-UT] skip failing tests

gulsumgudukbay committed 151 days ago

eae7c7cc

Add missing header `warnings` for test_util.py (#458)

psanal35 committed 151 days ago

8aa0e47b

Fix get_rocm_version to skip tests if installation is missing

psanal35 committed 151 days ago

f8f9fc20

[0.6.0-UT] Fix NumpyLinalgTest test cases (#453)

phambinhfin committed 151 days ago

8fc7e2b7

[0.6.0-UT] Add shared utility function get_rocm_version to test_util.py

psanal35 committed 151 days ago

871bec24

unskip test_select_n as it is passing on 0.6.0

gulsumgudukbay committed 151 days ago

29519680

[JAX] Fix numerical instability in testEighTinyNorm for ROCm GPU

phambinhfin committed 151 days ago

bf5ab2be

[ROCM] Add missing dependency on amdhip64 runtime

zahiqbal committed 155 days ago

9ecb94dc

git commit -m "Prepare for JAX release 0.7.0"

MichaelHudgins committed 164 days ago

8bc0ae6d

Rollback XLA archival change, current automation does not calculate the shasum correctly

MichaelHudgins committed 164 days ago

25a09917

[Mosaic GPU][NFC] Implement `inline_mgpu` and `custom_primitive` for WG semantics.

dimitar-asenov committed 164 days ago

295b92b9

Add visibility to `//jax:pretty_printer`

Jake VanderPlas committed 164 days ago

603d27ea

Update XLA dependency to use revision

Google-ML-Automation committed 165 days ago

c8e2fb20

Fix jax caching bug. The memory space selection logic should take into account the memory space of the requested sharding.

pschuh committed 165 days ago

b6e71051

Merge pull request #30368 from froystig:layout-doc

Google-ML-Automation committed 165 days ago

d16bf37c

write docs about controlling array layouts

froystig committed 165 days ago

641a768b

Remove `SafeNumDevices` now that all shardings are implemented in C++

yashk2810 committed 165 days ago

3eaa7f90

Loosen test tolerance for `tests/linalg_sharding_test.py::LinalgShardingTest::test_batch_axis_sharding_jvp`.

ZixuanJiang committed 165 days ago

2616b4d3

Merge pull request #29248 from DanisNone:ufuncs

Google-ML-Automation committed 165 days ago

36ae9d1b

Fix mesh_cast + new rng keys usage

yashk2810 committed 165 days ago

5371d6b7

[Pallas] Only run TPU interpret mode tests on CPU.

jburnim committed 165 days ago

8e21ca66

If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test.

yashk2810 committed 165 days ago

bd32431a

Fix `py_import` dependencies.

Google-ML-Automation committed 165 days ago

433a433a

[Pallas][Mosaic GPU] Fix racy test by having each warpgroup work on non-overlapping data.

Rifur13 committed 165 days ago

0562069c

[JAX] Purge more caches when clearing backends

hyeontaek committed 165 days ago

ffd90dae

Older