ROCm/jax

Pull Requests Commits

[ROCm]: call rocm function when lowering sparse ops

Rahul Batra committed 2 years ago

c95381d5

[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

add `threefry_partitionable` config setting to thread-local JIT context

froystig committed 2 years ago

b22e7571

Merge pull request #18269 from jakevdp:new-tutorials

a-googler committed 2 years ago

29c6faa0

Update XLA dependency to use revision

a-googler committed 2 years ago

57e33dc3

Merge pull request #18158 from mattjj:gpu-performance-docs

a-googler committed 2 years ago

002e2e79

draft docs on gpu performance tuning

mattjj committed 2 years ago

664e8347

Delete cached_call_jaxpr_lowerings since a more general cached_primitive_lowerings is available

yashk2810 committed 2 years ago

20255dce

[Try again] For nested pjit's cache the generation of StableHLO if it satifies the key. This should help in improving the lowering time.

yashk2810 committed 2 years ago

85af862e

Documentation: add stub files for new tutorial structure

jakevdp committed 2 years ago

344e42ba

Merge pull request #18324 from jakevdp:einsum-doc

a-googler committed 2 years ago

7eddc76a

Older