ROCm/jax

Pull Requests Commits

[JAX] [XLA:Python] Migrate pytree module to JAX.

hawkinsp committed 1 year ago

51560bf3

Add the missing `flatbuffers` dependency for the tests that run under `:build_jaxlib=false`.

Google-ML-Automation committed 1 year ago

49aad1b9

Update the docstring to mesh to use computation follows data and jax.jit APIs. Fixes https://github.com/jax-ml/jax/issues/27390

yashk2810 committed 1 year ago

c1904dc7

Skip checking of manylinux compliance for `jax` wheel.

Google-ML-Automation committed 1 year ago

5f1ab2ee

Reduced sharding in various tests.

mwhittaker committed 1 year ago

60b3e515

[Mosaic GPU] Adding pallas bindings to broadcast over the leading dimension and load a ref into WGMMAColFragLayout format.

Rifur13 committed 1 year ago

777d8f27

[pallas] Index Pallas refs instead of using `pl.load` and `pl.store`

superbobry committed 1 year ago

e7522639

[Mosaic GPU] Adding a new layout WGMMAColFragLayout to be able to load a 1d array and broadcast it along the leading dimension to a 2d shape as an input to a wgmma.

Rifur13 committed 1 year ago

ff718862

Merge pull request #27376 from jakevdp:jtu-type-annotations

Google-ML-Automation committed 1 year ago

89c7403d

Merge pull request #27363 from hawkinsp:pp

Google-ML-Automation committed 1 year ago

24d76ee1

Migrate more modules of xla/python to jax.

hawkinsp committed 1 year ago

9f3eb3e2

Increased tolerance in failing xla client tests.

mwhittaker committed 1 year ago

13b6e01a

Enable `jax` wheel testing via Bazel.

Google-ML-Automation committed 1 year ago

f5a4d1a8

jax.test_util: improve type annotations

jakevdp committed 1 year ago

7e235e3a

Merge pull request #27380 from hawkinsp:macbuild

Google-ML-Automation committed 1 year ago

ccc89652

Small cleanup to pretty-printer.

hawkinsp committed 1 year ago

13862ec1

Create `_FMA_FLOPS_FACTOR` to be used in roofline `dot` (and later `conv)`.

zacmustin committed 1 year ago

7e425396

Fix mac wheel build.

hawkinsp committed 1 year ago

94846941

Stop using mesh and `*_specs` in roofline tests.

zacmustin committed 1 year ago

b89cf0de

Delay the unflattening in `jnp.array`

superbobry committed 1 year ago

92f231e8

[Mosaic GPU] Adding a primitive to load from memrefs *with* a specified layout.

Rifur13 committed 1 year ago

a2f22cc1

Merge pull request #27266 from jakevdp:jax-array-test

Google-ML-Automation committed 1 year ago

315816d2

Merge pull request #26775 from ROCm:rocm-fix-numalib

Google-ML-Automation committed 1 year ago

014cf308

[pallas] Add support for `split` into any power-of-two equal parts in Triton lowering.

chr1sj0nes committed 1 year ago

198d7bb9

Update CUDA version in Bazel configs to 12.8, and CUDNN version to 9.8.

Google-ML-Automation committed 1 year ago

c1f65c3e

Change `python-tag` to `python_tag` to conform to the new setuptools version.

Google-ML-Automation committed 1 year ago

788ad8c6

Merge pull request #27364 from hawkinsp:tsan

Google-ML-Automation committed 1 year ago

d3836f16

Merge pull request #27337 from jburnim:jburnim_interpret_mode7

Google-ML-Automation committed 1 year ago

43ae8d70

[Mosaic GPU][NFC] Fix documentation of `WGMMA_LAYOUT`.

bchetioui committed 1 year ago

c6525bc5

Reenable tsan suppression, mark some tests as thread-unsafe.

hawkinsp committed 1 year ago

381f1109

Older