ROCm/jax

Pull Requests Commits

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 31 days ago

62f1e56c

Add download-jax-rocm-wheels composite action

mminutoli committed 31 days ago

7f3105e7

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 31 days ago

1cad7000

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 31 days ago

be9a2c2f

Improve OverFlow error raised by jit

yashk2810 committed 31 days ago

b9fdde25

Fix `get_sharding_for_vmap` to preserve unreduced

yashk2810 committed 31 days ago

14ce7f28

adding run_pytest_rocm.sh

gulsumgudukbay committed 31 days ago

9c629b47

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

yashk2810 committed 31 days ago

57a4f6a0

Propagate memory spaces from inputs to outputs when aliased.

brianwa84 committed 31 days ago

5f5be6f1

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

nitins17 committed 31 days ago

41256530

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

Google-ML-Automation committed 31 days ago

b2c727e3

Merge pull request #34770 from ROCm:add_support_of_building_jax_under_umbrella_workspace_new2

Google-ML-Automation committed 31 days ago

68992a18

Add missing bzlmod depenency

alekstheod committed 31 days ago

c6ad61b8

Add support to dynamically set test deps for jax as an external repo

alekstheod committed 31 days ago

1f4c34f2

Address review comments

alekstheod committed 31 days ago

60f50910

[mosaic] Added a new tpu.fetch_and_add_sync op

superbobry committed 31 days ago

2f1645f1

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

alekstheod committed 31 days ago

034c430c

Add some more information to deprecation messages in pxla modules

yashk2810 committed 31 days ago

0775e2b3

Reverts 4a5ad80ec771ec08bea2b28169d2f298aecbb46e

hawkinsp committed 31 days ago

5c4c7f05

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

Google-ML-Automation committed 31 days ago

8288263b

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

bchetioui committed 31 days ago

40c3359e

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

bchetioui committed 32 days ago

a5854e44

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

Google-ML-Automation committed 32 days ago

1066aa7b

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

superbobry committed 32 days ago

150a6d9d

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

superbobry committed 32 days ago

79fd127f

Regenerate python stubs for nanobind 2.11.0.

hawkinsp committed 32 days ago

6fb3f3da

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

Google-ML-Automation committed 33 days ago

3a93274b

Merge pull request #34462 from alekstheod:add_support_of_building_jax_under_umbrella_workspace

Google-ML-Automation committed 33 days ago

4a5ad80e

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

loislo committed 33 days ago

50b394dc

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

Google-ML-Automation committed 34 days ago

91871a92

Newer Older