ROCm/jax

Pull Requests Commits

Add missing wheel deps

alekstheod committed 49 days ago

ae98baa5

Merge pull request #39441 from mattjj:push-vvrkwosnukqt

Google-ML-Automation committed 50 days ago

56d7344d

Delete compute_on as compute_on2 is the replacement.

yashk2810 committed 50 days ago

3a440c76

[hijax] add test to verify custom_vjp works with hi types

mattjj committed 50 days ago

67aa9599

Don't call `update` on tokens in compute_on_abstract_eval

yashk2810 committed 50 days ago

dd430b64

Add partial_eval as the dep for mgpu/primitives.py

yashk2810 committed 50 days ago

727ce60b

Merge pull request #39569 from mattjj:push-zzntzkrwytvu

Google-ML-Automation committed 50 days ago

c83a51bd

eval_jaxpr_p should round-trip to eval_jaxpr_p

mattjj committed 50 days ago

5e95440f

Merge pull request #39508 from kodlan:fix/issue-7222-slice-index-error-message

Google-ML-Automation committed 50 days ago

5c4c6925

Only use /dev/shm/bazel output_base for TPU presubmit test jobs.

hawkinsp committed 51 days ago

51363ce6

Merge pull request #39565 from mattjj:push-rytmsxuwukls

Google-ML-Automation committed 51 days ago

ed0aad80

Make the default AxisType of `AbstractMesh`, `Explicit` instead of `Auto`.

yashk2810 committed 51 days ago

88f7a03b

[custom_vjp3] use eval_jaxpr_p to reduce trace times

mattjj committed 51 days ago

3ad6b053

Merge pull request #39564 from mattjj:push-suuxlnlqlpkq

Google-ML-Automation committed 51 days ago

8c231af6

[remat3] fixes for offloading, update tests

mattjj committed 51 days ago

e30c3120

[remat3] use eval_jaxpr_p in RematTraced.expand to reduce trace times

mattjj committed 51 days ago

ae1ef100

Support scalar prefetch in custom fusion.

IvyZX committed 51 days ago

1ca1aaf7

Add mirror rewrite for pythonhosted.org in JAX OSS Bazel downloader.

Google-ML-Automation committed 51 days ago

f387f296

Remove oldest_supported_libtpu testing configuration.

hawkinsp committed 51 days ago

aaf50c6a

fix test_fusible_shard_map_jit_grad

Cristian Garcia committed 51 days ago

73dcfc14

Merge pull request #39550 from ROCm:build-rocm-artifacts-cpu

Google-ML-Automation committed 51 days ago

95da8946

Rename `axis_shapes` to `axis_sizes` in `make_mesh`. No one uses this as a kwarg so it's safe to do.

yashk2810 committed 51 days ago

d78c899e

[Mosaic GPU] Move GPU module loading and function resolution out of runtime.

allanrenucci committed 51 days ago

4d786dec

Add more presubmit TPU tests.

hawkinsp committed 51 days ago

3c165171

[Mosaic] Support multiple batch dimensions.

WindQAQ committed 51 days ago

cadc19aa

[CI] Skip Pallas test_fusible_shard_map_jit_grad on TPU7x.

belitskiy committed 51 days ago

3eb27d36

Canonicalize ShapeDtypeStruct dtypes at JAX API boundaries using check_and_canonicalize_user_dtype.

hawkinsp committed 51 days ago

5a1619fb

Use one MLIR context per Python thread, rather than one per computation.

hawkinsp committed 51 days ago

044efe82

Build ROCm artifacts on CPU runners

psanal35 committed 51 days ago

5893b011

Merge pull request #39220 from ROCm:add_rocm_deps_to_plugin_wheels

Google-ML-Automation committed 51 days ago

299c9023

Older