ROCm/jax

Pull Requests Commits

use rocm_rbe also for build

alekstheod committed 253 days ago

cd77ba32

Add platform info

charleshofer committed 253 days ago

d5816720

Add RBE support in config

charleshofer committed 253 days ago

428042f9

Add remote config

charleshofer committed 257 days ago

cbbbc0c3

Add run tests ci script

alekstheod committed 259 days ago

2b8fa6a6

Enable unit test runs with Bazel

charleshofer committed 259 days ago

75640f9a

applied rocm-jax jax build fixes patches (#541)

zahiqbal committed 279 days ago

Verified 10d3836e

Prepare for JAX release 0.7.1

jakeharmon8 committed 314 days ago

5712de44

[Mosaic GPU] Separate hardware-agnostic from hardware-specific passes in BUILD files.

bchetioui committed 314 days ago

62482050

[Mosaic GPU] Annotate alignment of GMEM pointers to be 256

apaszke committed 314 days ago

a44869e1

[Mosaic GPU] Second attempt at fixing `test_tcgen05_collective_mma`.

allanrenucci committed 314 days ago

16371553

[Mosaic GPU] Attempt no. 2 at fixing the Mac OS build.

bchetioui committed 314 days ago

ffc0afb0

[Mosaic GPU] Do not predicate collective ops in the lowering.

allanrenucci committed 315 days ago

7ad15086

[Mosaic GPU] Fold non-gather indices into the gather TMA column index

apaszke committed 315 days ago

cbd0844e

Update XLA dependency to use revision http://github.com/openxla/xla/commit/31eb6029e8973337ef6c99810c27e1d807791c11

Google-ML-Automation committed 315 days ago

89628e45

[Mosaic GPU] Include `unistd.h` directly in order to fix build on Mac.

bchetioui committed 315 days ago

b09cc05f

Merge pull request #31151 from mattjj:fix-error-message

Google-ML-Automation committed 315 days ago

e44e8e65

fix a scan tree mismatch error

mattjj committed 315 days ago

cab4a539

[Mosaic:TPU] Mark tpu.relayout as pure with vector input/output, and clean up verifier and apply-vector-layout rule

tlongeri committed 315 days ago

f9dcd81b

Fix splash attention for non-128 divisible v head_dim

rdyro committed 315 days ago

c9a5799a

[Mosaic:TPU][Relayouts] Fix bug after cl/794807848

tlongeri committed 315 days ago

14c77fb1

Merge pull request #31140 from jakevdp:einsum-error

Google-ML-Automation committed 315 days ago

b438f851

Improve the error you get when you fail to construct an offloading policy. Or if the policy returns anything that's not a bool or `Recompute`, `Saveable` or `Offloable(...)`.

yashk2810 committed 315 days ago

d61c10e6

Merge pull request #31142 from jakevdp:ad-comment

Google-ML-Automation committed 315 days ago

1cac6f98

Fix incorrect comment in ad.py

jakevdp committed 315 days ago

335f3652

jnp.einsum: improve error for unsupported optimizations

jakevdp committed 315 days ago

d405ce2a

Merge pull request #31133 from jakevdp:polygamma-doc

Google-ML-Automation committed 315 days ago

31695257

[doc] typo in polygamma docs

jakevdp committed 315 days ago

0928d44a

Rollback due to downstream test failures

Google-ML-Automation committed 315 days ago

8d2c8ad5

Add `AbstractDevice` to `AbstractMesh` to capture specific device properties; `device_kind` and `num_tpu_cores`.

yashk2810 committed 315 days ago

8be91cde

Older