ROCm/jax

Pull Requests Commits

Add gfx942 to the list of supported devices

yuychang committed 2 years ago

8fe94175

Implement workaround for GPU-related error handling in prng_kernels.cc

yuychang committed 2 years ago

3bf8b679

Enhance GPU platform support for MI300 series

yuychang committed 2 years ago

234ab1d8

Add gfx940 and gfx941 to the list of supported devices

yuychang committed 2 years ago

c1f9f31c

[ROCm]: Disable some export tests

Rahul Batra committed 2 years ago

bb742a59

[ROCm]: Comment out GPU_LOWERING_ENABLED in sparse_test

Rahul Batra committed 2 years ago

da6b7d5d

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

75c05fb4

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

a9b6f30b

[ROCm]: Add jax-triton support for ROCm

Rahul Batra committed 2 years ago

59b1bea3

[ROCm]: Update build script

Rahul Batra committed 2 years ago

4bccaf56

Merge pull request #17989 from skye:version

a-googler committed 2 years ago

e27af485

Update setup.py and CHANGELOG for jax 0.4.18 release

skye committed 2 years ago

d4a1bb92

Merge pull request #17987 from jakevdp:lax-dep

a-googler committed 2 years ago

f0e4ea23

Merge pull request #17962 from jakevdp:fix-bitwise-count

a-googler committed 2 years ago

6681e64c

Merge pull request #17959 from jakevdp:lax-abs

a-googler committed 2 years ago

fae53d95

Merge pull request #17988 from froystig:rng-spec-jex

a-googler committed 2 years ago

5ac2513b

Merge pull request #17981 from hawkinsp:docs2

a-googler committed 2 years ago

31b680bb

Improve CUDA install documentation.

hawkinsp committed 2 years ago

8c4d020d

jax.lax: deprecate inadvertent exports & internal utilities

jakevdp committed 2 years ago

ce6a0c43

test custom PRNG impl construction round trip

froystig committed 2 years ago

7b831ba8

Merge pull request #17979 from hawkinsp:regex

a-googler committed 2 years ago

bd50c610

Merge pull request #17978 from hawkinsp:fft

a-googler committed 2 years ago

2052673f

Merge pull request #17983 from froystig:rng-spec

a-googler committed 2 years ago

6dee6f2c

identify PRNG schemes on key arrays, and recognize them in key constructors

froystig committed 2 years ago

5158e251

Merge pull request #17971 from hawkinsp:tanarm

a-googler committed 2 years ago

4c37c792

Merge pull request #17968 from jakevdp:random-test-util

a-googler committed 2 years ago

19b900d8

Add support for tiled layouts in arith.extf + implement (2,128)->(8,128) relayout for 32-bit types

apaszke committed 2 years ago

a703ec3d

[Mosaic] Strengthen the checks in the tpu.iota apply_vector_layout rule

apaszke committed 2 years ago

327359cb

Fix NumPy warning when casting negative float to unsigned in tan() test on aarch64 under Python 3.12

hawkinsp committed 2 years ago

5a271212

Fix tests that mistakenly used assertRaises(..., msg=...) to match a message.

hawkinsp committed 2 years ago

53117468

Older