ROCm/jax

Pull Requests Commits

Lower block size of attention op to 32

mqhc2020 committed 2 years ago

50e25c15

Extract hsaco path from compiled kernel tuple

mqhc2020 committed 2 years ago

7965e1b2

[ROCm]: Update docker file for ROCm6.0w

Rahul Batra committed 2 years ago

ee929e56

Set blk size to 64 for attn opn and disable some pallas tests

Rahul Batra committed 2 years ago

933702a5

[ROCm]: Add ROCm Target

Rahul Batra committed 2 years ago

c99438d5

[ROCm]: Updates for ROCm6.0

Rahul Batra committed 2 years ago

172141f5

Add gfx942 to the list of supported devices

yuychang committed 2 years ago

70ae2a32

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

yuychang committed 2 years ago

599e55b5

Enhance GPU platform support for MI300 series

yuychang committed 2 years ago

d9789b6c

Add gfx940 and gfx941 to the list of supported devices

yuychang committed 2 years ago

4a8cfcae

[ROCm]: Disable some export tests

Rahul Batra committed 2 years ago

1a4f5ff8

[ROCm]: Comment out GPU_LOWERING_ENABLED in sparse_test

Rahul Batra committed 2 years ago

ef68c4ff

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

e6a04026

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

19da4c01

[ROCm]: Update build script

Rahul Batra committed 2 years ago

9fa53699

Merge pull request #18325 from skye:version

a-googler committed 2 years ago

e22b8bbc

Update versions for jax + jaxlib 0.4.20 release

skye committed 2 years ago

68138191

Fix typing annotations for `@jax.named_call`

Conchylicultor committed 2 years ago

81ac67f3

Update XLA dependency to use revision

a-googler committed 2 years ago

1c66ac53

Properly pack and unpack int4 arrays on CPU in PJRT.

reedwm committed 2 years ago

d41078fb

[Mosaic] Add support for generalized reduction unrolling to infer_vector_layout

apaszke committed 2 years ago

5d289615

Add support for dynamic indices in VMEM loads and stores

apaszke committed 2 years ago

f17f5492

Update XLA dependency to use revision

a-googler committed 2 years ago

32a317f7

[Mosaic][NFC] Remove unused parameter

tlongeri committed 2 years ago

caa1c39a

Pass flags from kernel into HLO backend config.

a-googler committed 2 years ago

a009f8d6

Merge pull request #18340 from jakevdp:keyarray-error

a-googler committed 2 years ago

9f28512c

[jex] replace `extend.random.PRNGImpl` with `extend.random.define_prng_impl`

froystig committed 2 years ago

16d082b0

[random] add more information to KeyArray deprecation error

jakevdp committed 2 years ago

a4e6b4e9

Merge pull request #18338 from froystig:partitionable-threefry-ctx-mgr

a-googler committed 2 years ago

49fedb1c

add `jax.threefry_partitionable` context manager

froystig committed 2 years ago

ed9a4c29

Older