ROCm/jax

Pull Requests Commits

[ROCm] rnn implementation

Ruturaj4 committed 1 year ago

0eefd958

Ruturaj4 committed 1 year ago

6001ed39

[ROCm] jaxlib linalg fix

Ruturaj4 committed 1 year ago

1c75f5f9

Ruturaj4 committed 1 year ago

241152c7

Fix DOCKER_BUILDKIT issue for jax CI

kiran-thumma committed 1 year ago

9ca504a3

Fix formatting issues

mrodden committed 1 year ago

8b11832d

[ROCm] Fix invalid repo url for EL path

mrodden committed 1 year ago

8ff96792

Change run_multi_gpu set opts

mrodden committed 1 year ago

83e3d0a8

BUILD file fixes and linter run

Ruturaj4 committed 1 year ago

9be19a7e

[ROCm] Bring up clang support for JAX+XLA

Ruturaj4 committed 1 year ago

dfb7db0e

[ROCm] Fix formatting on python files

mrodden committed 1 year ago

9dbbb3a3

[ROCm] Use specific amdgpu version for EL8 systems

mrodden committed 1 year ago

b2c1b497

[ROCm] Remove broken legacy env vars

mrodden committed 1 year ago

f27e9d61

[ROCM] fix typename

Ruturaj4 committed 1 year ago

aff7f1e4

Simplify logic in jaxlib FFI_ASSIGN_OR_RETURN macro, and fix gcc build.

dfm committed 1 year ago

df62e8db

Prepare for v0.4.33 release.

hawkinsp committed 1 year ago

80e1c94d

Prepare for v0.4.32 release.

hawkinsp committed 1 year ago

1594d2f3

Make sure to call the superclass' __init__() on a newly created instance in PositionalSharding._remake().

hawkinsp committed 1 year ago

ed849ff9

Relax test tolerance in pinv test to fix a CI failure on Windows CPU.

hawkinsp committed 1 year ago

2bd1fdea

Merge pull request #23415 from kaixih:key_value_seq_lengths

Google-ML-Automation committed 1 year ago

e869a9d6

Internal change

superbobry committed 1 year ago

ea68f456

Disable a pallas export compatibility test that fails on TPU v6e.

hawkinsp committed 1 year ago

49dd6ed8

Update users of jax.tree.map() to be more careful about how they handle Nones.

hawkinsp committed 1 year ago

808003b4

[Pallas] Implement tiled and swizzled Memref loads for Mosaic GPU via "GPUBlockSpec"

justinjfu committed 1 year ago

e3c4b20f

[Pallas] Disable win32 gpu_ops_test.

justinjfu committed 1 year ago

c659dc9a

Merge pull request #23549 from pschuh:docs-update

a-googler committed 1 year ago

14b86259

Reverts 5e4250e64bb415be94ddc8a80dba6083a6a4123a

yashk2810 committed 1 year ago

0b04dd02

Removed a sneaky comma in Pallas Mosaic GPU lowering

superbobry committed 1 year ago

e7b261c3

Prepare for jax 0.4.32 release

yashk2810 committed 1 year ago

5e4250e6

Clean up JAX backend for all backends to avoid dangling PyClient references.

Google-ML-Automation committed 1 year ago

4957ab9a

Older