ROCm/jax

Pull Requests Commits

Add FA bwd split kernels

Rahul Batra committed 2 years ago

ab66a48b

Add support for transposing tensor being loaded

vgokhale committed 2 years ago

25d62ff5

Convert P to bf16 instead of f16

vgokhale committed 2 years ago

4547ff96

Add support for exp2

vgokhale committed 2 years ago

23c7c46c

Update FA to v2

vgokhale committed 2 years ago

824f36f9

[ROCm]: Don't set dynamic shmem on ROCM for triton

Rahul Batra committed 2 years ago

fa9aba68

[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

Older