ROCm/jax

Pull Requests Commits

[Fix] Nvidia wheel verion arg should not be required

GeneDer committed 319 days ago

Verified 92af834b

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

zahiqbal committed 334 days ago

Verified 10d3836e

Prepare for JAX release 0.7.1

jakeharmon8 committed 1 year ago

5712de44

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

bchetioui committed 1 year ago

62482050

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

apaszke committed 1 year ago

a44869e1

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

allanrenucci committed 1 year ago

16371553

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

bchetioui committed 1 year ago

ffc0afb0

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

allanrenucci committed 1 year ago

7ad15086

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

apaszke committed 1 year ago

cbd0844e

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

Google-ML-Automation committed 1 year ago

89628e45

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

bchetioui committed 1 year ago

b09cc05f

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

Google-ML-Automation committed 1 year ago

e44e8e65

fix a scan tree mismatch error

mattjj committed 1 year ago

cab4a539

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

tlongeri committed 1 year ago

f9dcd81b

Fix splash attention for non-128 divisible v head_dim

rdyro committed 1 year ago

c9a5799a

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

tlongeri committed 1 year ago

14c77fb1

Merge pull request #31140 from jakevdp:einsum-error

Google-ML-Automation committed 1 year 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 1 year ago

d61c10e6

Merge pull request #31142 from jakevdp:ad-comment

Google-ML-Automation committed 1 year ago

1cac6f98

Fix incorrect comment in ad.py

jakevdp committed 1 year ago

335f3652

jnp.einsum: improve error for unsupported optimizations

jakevdp committed 1 year ago

d405ce2a

Merge pull request #31133 from jakevdp:polygamma-doc

Google-ML-Automation committed 1 year ago

31695257

[doc] typo in polygamma docs

jakevdp committed 1 year ago

0928d44a

Rollback due to downstream test failures

Google-ML-Automation committed 1 year ago

8d2c8ad5

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

yashk2810 committed 1 year ago

8be91cde

Merge pull request #31131 from mattjj:discharge-state-caching

Google-ML-Automation committed 1 year ago

5977fc2e

[mutable-arrays] add discharge_state2 for ClosedJaxprs, with caching

mattjj committed 1 year ago

5ff55080

Expose a standard function to insert pvary in `jax.random` APIs that we can reuse.

yashk2810 committed 1 year ago

9d1661ed

[Mosaic GPU] Use a custom pass to compile PTX to CUBIN.

bchetioui committed 1 year ago

87d5de4f

Make jax.lax.empty support uninitialized values on TPU and expose the API as `jax.lax.empty`

yashk2810 committed 1 year ago

99d4cc89

Older