ROCm/jax

Pull Requests Commits

[ROCm] Add float8_e4m3fnuz and float8_e5m2fnuz support for Rocm

Ruturaj4 committed 2 years ago

41e4c25d

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

3f1b0595

Refactor & test internal deprecation APIs

Jake VanderPlas committed 2 years ago

d33a5689

Merge pull request #21319 from gnecula:exp_fix_mesh

a-googler committed 2 years ago

2eff241f

[export] Fix device assignment error for grad of exported.

gnecula committed 2 years ago

6deeee27

[Mosaic] Also check bitwidth in apply-vector-layout's `layoutIsValidForValue`.

tlongeri committed 2 years ago

b197ae52

Merge pull request #21318 from jakevdp:fix-mypy

a-googler committed 2 years ago

118ca21b

CI: fix mypy error

jakevdp committed 2 years ago

329ab036

Copybara import of the project:

Cjkkkk committed 2 years ago

06d2e489

Simplify JAX lowering rules for cumulative sum

cheshire committed 2 years ago

e0a6453a

Add unit test to check if the backend serialization/deserialization result equal to the original executable.

Google-ML-Automation committed 2 years ago

cc3a380f

Add support for TPU delay in Mosaic

Google-ML-Automation committed 2 years ago

61ff8287

[Mosaic GPU] Prepare matmul example so it can be exposed to other projects.

Google-ML-Automation committed 2 years ago

2f45830d

Use @register_lowering to register Pallas GPU lowering rules

superbobry committed 2 years ago

f600caa2

Merge pull request #21305 from jakevdp:scalar-bool

a-googler committed 2 years ago

c3d0b0d1

Merge pull request #21292 from ROCm:rv_stable_051624

a-googler committed 2 years ago

974c72b9

Finalize deprecation of the config module.

Jake VanderPlas committed 2 years ago

4bac10e7

Merge pull request #21231 from nouiz:doc_experimental_serialize_executable

a-googler committed 2 years ago

bb616eff

Support scalar boolean indices in arr.at[idx].set(vals)

jakevdp committed 2 years ago

5b28170b

Add notap tag to Mosaic tests

apaszke committed 2 years ago

53ec2cd2

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

ffdb9bb0

fix: Update hermetic python dependencies to numpy=2.0.0rc2 and scipy=1.13.0 for all python version

vam-google committed 2 years ago

45a7c22e

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

8caeaa29

Make `eqn.ctx` context manager thread safe by creating `eqn.ctx.manager`.

yashk2810 committed 2 years ago

6577f47b

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

e3a7a87f

Support remat + compute_on. If the rematted computation is annotated to run on host, the backward pass will also execute on host. Also enable no-op nested compute tests.

yashk2810 committed 2 years ago

25aa13c4

add cholesky changes in bazel

Ruturaj4 committed 2 years ago

79fccf6c

jax/pallas support ellipsis indexing

Google-ML-Automation committed 2 years ago

641d5c8b

Make `jax.grad` and `compute_on` work correctly. If the forward pass has annotation to execute on CPU, then it's backward pass also executes on CPU.

yashk2810 committed 2 years ago

02c19e96

Add quantization support for PagedAttention TPU Pallas kernel.

ashishenoyp committed 2 years ago

1043e24f

Older