ROCm/jax

Pull Requests Commits

[ROCm]: Add jax-triton support for ROCm

Rahul Batra committed 2 years ago

310debe8

[ROCm]: Add hipblaslt in Dockerfile

Rahul Batra committed 2 years ago

08eb9b96

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

fb822101

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

3bbeabe9

Merge pull request #16527 from skye:version

a-googler committed 2 years ago

c3e24270

Jax 0.4.13 release.

hawkinsp committed 2 years ago

487b640a

Update JAX's XlaExecutable.cost_analysis and related plumbing so it works on Cloud TPU

skye committed 2 years ago

10424c59

Silence pytype errors under an upcoming pytype change.

hawkinsp committed 2 years ago

9f4080ae

Merge pull request #16508 from hawkinsp:metal

a-googler committed 2 years ago

e123d1e9

Ignore `JAX_USE_PJRT_C_API_ON_TPU=false` user warning raised.

yashk2810 committed 2 years ago

677b0d9d

Add a link to the Apple Metal plugin to the JAX README.

hawkinsp committed 2 years ago

85a84fde

Zlib compress kernel proto.

chr1sj0nes committed 2 years ago

b3527f39

Make JAX-Triton calls serializable.

chr1sj0nes committed 2 years ago

f2386674

[jax2tf] Refactor the backwards compatibility tests.

gnecula committed 2 years ago

c5a47d18

[shape_poly] linalg.cholesky: add support for shape polymorphism with native lowering on CPU

gnecula committed 2 years ago

f02d1223

[shape_poly] linalg.qr: shape polymorphism with native lowering on CPU

gnecula committed 2 years ago

d940db60

[shape_poly] linalg.svd: shape polymorphism for native serialization on CPU

gnecula committed 2 years ago

92288d30

Documentation improvements for XlaCallModule.disabled_checks.

gnecula committed 2 years ago

1535fa09

Merge pull request #16499 from skye:se_warning

a-googler committed 2 years ago

5742ab4d

Use `EncapsulateFunction` utility.

chr1sj0nes committed 2 years ago

64e73270

Fix `jax.clear_backends()` on Cloud TPU

skye committed 2 years ago

529822cb

Merge pull request #16484 from jakevdp:key-shape

a-googler committed 2 years ago

f46e5141

Merge pull request #16404 from jakevdp:tracer-errors

a-googler committed 2 years ago

06f76bc6

Merge pull request #16333 from skye:profiler_test

a-googler committed 2 years ago

961d9188

profiler_test.py fixes and add coverage to Cloud TPU CI

skye committed 2 years ago

2ca151ef

Issue warning if JAX_USE_PJRT_C_API_ON_TPU=false

skye committed 2 years ago

71cbfe49

Merge pull request #16498 from ROCmSoftwarePlatform:rocm-update-build-defaults

a-googler committed 2 years ago

3cfb57d4

Fix symbol collision when merging MLIR modules.

hawkinsp committed 2 years ago

e99ca460

Merge pull request #16497 from hawkinsp:workflows

a-googler committed 2 years ago

eca3b972

Rename Github actions so they have more consistent naming.

hawkinsp committed 2 years ago

741f0d03

Older