ROCm/jax

Pull Requests Commits

[pallas:mosaic_gpu] Add support for run_scoped

cperivol committed 2 years ago

4186824b

[XLA:Mosaic] Create apply layout pass with ctx instead of config list.

bythew3i committed 2 years ago

7e2107b1

[pallas:mosaic_gpu] Take into account that mosaic gpu now supports clusters.

cperivol committed 2 years ago

101e5fe1

Shmallas, a.k.a. allow lowering shard_map + run_state to a pallas_call.

sharadmv committed 2 years ago

ae8da833

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

5b0be26a

#sdy Add dataclass that stores the data needed to build a Shardy MLIR sharding attribute.

bartchr808 committed 2 years ago

18f4456f

[xla:cpu] Support for up to 16 sorted inputs

ezhulenev committed 2 years ago

e3fc63ca

Merge pull request #22597 from jakevdp:arr-device

a-googler committed 2 years ago

dc42ba0e

Merge pull request #22574 from jakevdp:xla-computation

a-googler committed 2 years ago

7792bded

Improved docs for jnp.fft.fft2 and ifft2

rajasekharporeddy committed 2 years ago

1525f012

[array API] add device property & to_device method

jakevdp committed 2 years ago

613a0004

Merge pull request #22594 from jakevdp:compress-pyi

a-googler committed 2 years ago

13e42ad4

Merge pull request #22588 from rajasekharporeddy:testbranch3

a-googler committed 2 years ago

ac4a7b1e

Improved docs for jnp.tri, tril and triu

rajasekharporeddy committed 2 years ago

f90d0ee0

Merge pull request #22487 from rajasekharporeddy:testbranch1

a-googler committed 2 years ago

0264322a

Skip some Pallas tests that fail on TPUv6.

hawkinsp committed 2 years ago

fd85c783

Add missing parameters to jnp.compress type interface

jakevdp committed 2 years ago

a88a4b13

Merge pull request #22559 from superbobry:pallas-test

a-googler committed 2 years ago

0c09e794

Merge pull request #22558 from superbobry:pallas

a-googler committed 2 years ago

696e7304

#sdy Initial set of changes to allow for lowering to the Shardy dialect.

bartchr808 committed 2 years ago

864178d3

Remove the unaccelerate_deprecation utility

jakevdp committed 2 years ago

f887b66d

Reverts 093b92be8ed7bd979486614325956e88cc474ff1

gnecula committed 2 years ago

459b83cf

[Mosaic GPU] Add a collective mbarrier interface

apaszke committed 2 years ago

f0792b2d

Another take at enabling Pallas GPU tests on x64

superbobry committed 2 years ago

b7715e27

[MosaicGPU] Add a __repr__ to FragmentedArray so pallas:mosaic_gpu errors are more readable.

cperivol committed 2 years ago

81cb1add

[Mosaic GPU] Replace multicast_mask by a nicer collective async copy interface

apaszke committed 2 years ago

51732c5c

[Mosaic GPU] Add early support for block clusters and multicast TMA

apaszke committed 2 years ago

a2b2fbf5

Add support for named grids in pallas_call.

sharadmv committed 2 years ago

499ceeeb

Improved docs for jnp.fft.rfft and irfft

rajasekharporeddy committed 2 years ago

1650d1e8

Reverts d7b821b04d8fec543f570faaece7572a50a75eb6

gnecula committed 2 years ago

a18872aa

Newer Older