ROCm/jax

Pull Requests Commits

Disable lobpcg test

Rahul Batra committed 2 years ago

21a066e5

Update docker file to ROCm 5.5

Rahul Batra committed 2 years ago

5f5cfa53

Update docker file to ROCm 5.5

Rahul Batra committed 2 years ago

ec571612

[ROCm]: Run pmap test on specific number of GPUs

invalid-email-address committed 3 years ago

748c93a5

[ROCm]: Update to ci script

Rahul Batra committed 3 years ago

7b1f90f4

Hardcode runtime flag

invalid-email-address committed 3 years ago

9251cd5a

[ROCm]: Disable api_test remat checkpoint dots outside scan

invalid-email-address committed 3 years ago

4d016d83

[ROCm]: Disable some scipy stats test

invalid-email-address committed 3 years ago

fa27a864

[ROCm]: Disable sparse tests

invalid-email-address committed 3 years ago

475bb6d1

[ROCm]: Disable laxNumpy testSelect

invalid-email-address committed 3 years ago

1a532f5f

[ROCm]: Remove bfloat16 dtype from unit tests

invalid-email-address committed 3 years ago

cda7f4b9

[ROCm]: Disable linalg_test::SlogdetGrad

invalid-email-address committed 3 years ago

57f88d94

Merge remote-tracking branch 'rocm/rocm-jaxlib-v0.4.6' into rocm-jaxlib-v0.4.6

Rahul Batra committed 3 years ago

000d57a9

[ROCm]: Updates for container and build script

Rahul Batra committed 3 years ago

8018aa8b

[ROCm] Disable some compilation cache tests

Rahul Batra committed 3 years ago

1c039498

[ROCm]: Add support for ROCm jaxlib wheel

Rahul Batra committed 3 years ago

57c091cf

[ROCm]: Updates for container and build script

Rahul Batra committed 3 years ago

3cd9739d

[ROCm]: Add support for ROCm jaxlib wheel

Rahul Batra committed 3 years ago

ce5074c2

Merge pull request #14881 from skye:version

a-googler committed 3 years ago

f1f4840a

Update WORKSPACE and setup.py for jax/jaxlib 0.4.6 release

skye committed 3 years ago

1aa08fd4

Split _src/traceback_util.py into its own Bazel target.

hawkinsp committed 3 years ago

7fd1e2ff

Increase sharding of shard_map test on CPU.

hawkinsp committed 3 years ago

01b00c48

Split _src/pretty_printer.py into its own Bazel target.

hawkinsp committed 3 years ago

9912a8eb

Remove the disassemble into single devices arrays in ExecuteReplicated.__call__ since it is not needed anymore since shard_args return Array only with `xla_extension_version >= 136`.

yashk2810 committed 3 years ago

d216d98c

Exclude "util.py" and "config.py" from the main JAX bazel target.

hawkinsp committed 3 years ago

08789fd9

[jax2tf] Disable some failing tests

gnecula committed 3 years ago

560fe73e

Split some submodules out of //jax under Bazel.

hawkinsp committed 3 years ago

0e05a798

[jax2tf] Add check that native lowering should not include custom calls not guaranteed to be stable.

gnecula committed 3 years ago

5c914534

Results handlers are now able to take Arrays.

pschuh committed 3 years ago

50d83583

Merge pull request #14863 from mattjj:shmap-debug-print

a-googler committed 3 years ago

3a9a1ba2

Older