ROCm/jax

Pull Requests Commits

Expand ROCm testing matrix and improve job naming

mminutoli committed 17 hours ago

f9d2e1b5

Add JAXCI_ROCM_VERSION environment variable

mminutoli committed 17 hours ago

6ad760e6

Add ROCm pytest workflow to nightly/release CI

mminutoli committed 17 hours ago

f952348c

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 17 hours ago

62f1e56c

Add download-jax-rocm-wheels composite action

mminutoli committed 17 hours ago

7f3105e7

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 17 hours ago

1cad7000

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 17 hours ago

be9a2c2f

adding run_pytest_rocm.sh

gulsumgudukbay committed 21 hours ago

9c629b47

Propagate memory spaces from inputs to outputs when aliased.

brianwa84 committed 1 day ago

5f5be6f1

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

nitins17 committed 1 day ago

41256530

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

Google-ML-Automation committed 1 day ago

b2c727e3

Merge pull request #34770 from ROCm:add_support_of_building_jax_under_umbrella_workspace_new2

Google-ML-Automation committed 1 day ago

68992a18

Address review comments

alekstheod committed 1 day ago

60f50910

[mosaic] Added a new tpu.fetch_and_add_sync op

superbobry committed 1 day ago

2f1645f1

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

alekstheod committed 1 day ago

034c430c

Add some more information to deprecation messages in pxla modules

yashk2810 committed 1 day ago

0775e2b3

Reverts 4a5ad80ec771ec08bea2b28169d2f298aecbb46e

hawkinsp committed 1 day ago

5c4c7f05

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

Google-ML-Automation committed 1 day ago

8288263b

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

bchetioui committed 1 day ago

40c3359e

[Pallas/Mosaic GPU] Remove unnecessary jaxlib version guard-related code.

bchetioui committed 1 day ago

a5854e44

Update XLA dependency to use revision http://github.com/openxla/xla/commit/80529bdfc1c3bc206d0809a474031e9c9008c9f0

Google-ML-Automation committed 1 day ago

1066aa7b

[mosaic] Added a macro which defines AbslStringify for all MLIR enums in the enclosing namespace

superbobry committed 2 days ago

150a6d9d

[mosaic] Avoid using the deprecated `output_memory_colors` in `CustomCallBackendConfig`

superbobry committed 2 days ago

79fd127f

Regenerate python stubs for nanobind 2.11.0.

hawkinsp committed 2 days ago

6fb3f3da

Update XLA dependency to use revision http://github.com/openxla/xla/commit/2a7188001603266f57a661a10beee3418412c312

Google-ML-Automation committed 2 days ago

3a93274b

Merge pull request #34462 from alekstheod:add_support_of_building_jax_under_umbrella_workspace

Google-ML-Automation committed 3 days ago

4a5ad80e

Adds `jax.lax.scaled_dot` for scaled dot products.

loislo committed 3 days ago

50b394dc

Update XLA dependency to use revision http://github.com/openxla/xla/commit/47fe49edce6c27225a0cb6675bf08d43550316a6

Google-ML-Automation committed 3 days ago

91871a92

[Mosaic GPU][NFC] Clean up the lowering of `slice_smem`.

bchetioui committed 3 days ago

01d76c43

Update jax.host_id() and jax.host_count() to jax.process_index() and jax.process_count(), respectively.

danielsuo committed 3 days ago

46c02566

Older