ROCm/jax

Pull Requests Commits

update gpu ram per test

gulsumgudukbay committed 154 days ago

07278c5b

remove JAX_ENABLE_CUDA_XDIST

gulsumgudukbay committed 154 days ago

22e3634e

fixing aborts caused by xdist

gulsumgudukbay committed 154 days ago

fc1b6a38

Add check for /dev/shm

mminutoli committed 154 days ago

284004bf

[TESTING ONLY] Bypass ROCm plugin version check and strip nightly workflow

mminutoli committed 154 days ago

a9515347

[lapack] Parallelize SVD computations across the batch dimension.

basioli-k committed 154 days ago

9aed01b0

Update XLA dependency to use revision http://github.com/openxla/xla/commit/2734172e27d22c1e774b4d27365074e09e4c1c6d

Google-ML-Automation committed 154 days ago

f60c75e1

Improve OverFlow error raised by jit

yashk2810 committed 154 days ago

cf122636

Fix `get_sharding_for_vmap` to preserve unreduced

yashk2810 committed 154 days ago

938242dd

Update transpose rules of all primitives that define an unreduced_rule to call `.to_cotangent_aval()` on residuals.

yashk2810 committed 154 days ago

8750ec7f

Standardize parameter naming across workflows and actions

mminutoli committed 154 days ago

675e05cb

Expand ROCm testing matrix and improve job naming

mminutoli committed 156 days ago

f9d2e1b5

Add JAXCI_ROCM_VERSION environment variable

mminutoli committed 156 days ago

6ad760e6

Add ROCm pytest workflow to nightly/release CI

mminutoli committed 156 days ago

f952348c

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 156 days ago

62f1e56c

Add download-jax-rocm-wheels composite action

mminutoli committed 156 days ago

7f3105e7

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 156 days ago

1cad7000

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 156 days ago

be9a2c2f

adding run_pytest_rocm.sh

gulsumgudukbay committed 156 days ago

9c629b47

Propagate memory spaces from inputs to outputs when aliased.

brianwa84 committed 156 days ago

5f5be6f1

Switch all Windows RBE builds to use `ml-oss-rbe-general`

nitins17 committed 156 days ago

41256530

Save CPU version of collective metadata to resources so it can be used during execution.

Google-ML-Automation committed 156 days ago

b2c727e3

Merge pull request #34770 from ROCm:add_support_of_building_jax_under_umbrella_workspace_new2

Google-ML-Automation committed 156 days ago

68992a18

Address review comments

alekstheod committed 156 days ago

60f50910

[mosaic] Added a new tpu.fetch_and_add_sync op

superbobry committed 156 days ago

2f1645f1

Add support of executing jax tests from withing a parent workspace as @jax//tests:gpu_tests

alekstheod committed 156 days ago

034c430c

Add some more information to deprecation messages in pxla modules

yashk2810 committed 156 days ago

0775e2b3

Reverts 4a5ad80ec771ec08bea2b28169d2f298aecbb46e

hawkinsp committed 156 days ago

5c4c7f05

Merge pull request #34682 from olupton:workaround-tcgen05-llvm-errors

Google-ML-Automation committed 156 days ago

8288263b

[Mosaic GPU] Clean up all the linter warnings in `dialect_lowering.py`.

bchetioui committed 156 days ago

40c3359e

Older