ROCm/jax

Pull Requests Commits

ROCM: Update dockerfile to ROCm5.6

Rahul Batra committed 2 years ago

fc46bbaa

[ROCm] Disable testCondGrad

Rahul Batra committed 2 years ago

afded4bb

[ROCm]: Disable laxNumpy testSelect

invalid-email-address committed 2 years ago

bac00ebb

[ROCm]: Disable some scipy stats test

invalid-email-address committed 2 years ago

df89bf40

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

843c9c2e

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

99994b86

Disable lobpcg test

Rahul Batra committed 2 years ago

f43c9fc0

[ROCm]: Remove bfloat16 dtype from unit tests

invalid-email-address committed 2 years ago

396dbac5

Switch to using OpenXLA

Rahul Batra committed 2 years ago

cd619ed6

Wheel name w/o ROCm, add Python lib dependencies

Rahul Batra committed 2 years ago

62ae5a9e

[ROCm]: Updates for container and build script

Rahul Batra committed 2 years ago

95c9d8cb

Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.9 release, take 2

skye committed 2 years ago

5bcd9dcc

Fix the cache on `to_gspmd_sharding` to depend on if device/backend is set on pjit/jit.

yashk2810 committed 2 years ago

befa29b5

Use set equality operators instead of intersection because I didn't know set had equality operators.

yashk2810 committed 2 years ago

2694bf62

Merge pull request #15929 from gnecula:fix_mlir_ir

a-googler committed 2 years ago

68ba5424

[GPU] Fix another instance of missing stream synchronization in RNN kernels.

hawkinsp committed 2 years ago

a89c3777

Merge pull request #15931 from geraschenko:bcoo_reshape

a-googler committed 2 years ago

a2b5bd52

Make `dimensions` argument of bcoo_reshape optional.

geraschenko committed 2 years ago

27aa5fb7

Add McJAX resharding to device_put. Allow resharding if inputs and target sharding have the same set of devices but different order.

yashk2810 committed 2 years ago

18d19caa

Merge pull request #15920 from froystig:issue15869

a-googler committed 2 years ago

cf4c1eda

Import "ir" directly rather than as "mlir.ir"

gnecula committed 2 years ago

daf6a30f

Merge pull request #15859 from gnecula:poly_rng

a-googler committed 2 years ago

cb3a4f3d

[shape_poly] Improvements and more testing for shape polymorphism for random primitives

gnecula committed 2 years ago

de2a811f

[shape_poly] Improve testing of vmap test harnesses

gnecula committed 2 years ago

16881e62

[shape_poly] Add support for shape polymorphism with native serialization for lax.linalg.qr on TPU

gnecula committed 2 years ago

d66be780

delegate `select` lowering to opaque dtype rule

froystig committed 2 years ago

051c5dda

Merge pull request #15909 from skye:version

a-googler committed 2 years ago

236c74ca

[GPU] Add missing stream synchronization to tridiagonal_solve gtsv2 call.

hawkinsp committed 2 years ago

f168a156

Add tests for negative inputs to top-k.

hawkinsp committed 2 years ago

00b75aff

Use stream-synchronized copy in rnn_kernels.cc.

hawkinsp committed 2 years ago

6b9a1099

Older