ROCm/jax

Pull Requests Commits

Let the unit tests use build.py for setting up Bazel commands for uni… (#583)

charleshofer committed 7 days ago

Verified da92504f

Fix print statement syntax for ROCm error message (#573)

gulsumgudukbay committed 11 days ago

Verified b14ed378

Remove old scripts and break build for ROCm (#568)

gulsumgudukbay committed 11 days ago

Verified 09c8b8cb

Fix v_pages quantization and adjust test params for ROCm compatibility (#559)

phambinhfin committed 32 days ago

Verified 00e65c32

Fix hipSparse CSR algorithm mappings for ROCm 7 (#557)

phambinhfin committed 34 days ago

Verified ccf941ee

Make nvidia version data optional for ROCm builds (#552)

phambinhfin committed 44 days ago

Verified fd3e2c6d

Phambinh/jax fix copy (#554)

phambinhfin committed 44 days ago

Verified 3206662a

Skip test_dot_precision() with TF32 for gfx11xx and gfx12xx architecture (#520) (#551)

AleksaArsic committed 55 days ago

Verified 3b02da3a

The ROCm plugin should fail gracefully and allow for other devices to… (#550)

JehandadKhan committed 60 days ago

Verified 0011fb97

Porting v0.6.0 changes to v0.7.1 (#549)

zahiqbal committed 74 days ago

Verified a5d5fa0d

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

zahiqbal committed 90 days ago

Verified 10d3836e

Prepare for JAX release 0.7.1

jakeharmon8 committed 125 days ago

5712de44

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

bchetioui committed 125 days ago

62482050

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

apaszke committed 125 days ago

a44869e1

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

allanrenucci committed 125 days ago

16371553

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

bchetioui committed 125 days ago

ffc0afb0

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

allanrenucci committed 125 days ago

7ad15086

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

apaszke committed 125 days ago

cbd0844e

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

Google-ML-Automation committed 125 days ago

89628e45

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

bchetioui committed 125 days ago

b09cc05f

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

Google-ML-Automation committed 125 days ago

e44e8e65

fix a scan tree mismatch error

mattjj committed 125 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 125 days ago

f9dcd81b

Fix splash attention for non-128 divisible v head_dim

rdyro committed 125 days ago

c9a5799a

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

tlongeri committed 125 days ago

14c77fb1

Merge pull request #31140 from jakevdp:einsum-error

Google-ML-Automation committed 125 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 125 days ago

d61c10e6

Merge pull request #31142 from jakevdp:ad-comment

Google-ML-Automation committed 125 days ago

1cac6f98

Fix incorrect comment in ad.py

jakevdp committed 125 days ago

335f3652

jnp.einsum: improve error for unsupported optimizations

jakevdp committed 125 days ago

d405ce2a

Older