ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

bb6a1fba

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b418ee41

[Mosaic] Fold load->reshape and reshape->store into strided load/store respectively.

Google-ML-Automation committed 1 year ago

ff0a516f

Merge pull request #30173 from gnecula:print_large_constants

Google-ML-Automation committed 1 year ago

f56cb0ee

Summarize large constants that may appear in Jaxpr literals

gnecula committed 1 year ago

3b386790

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c0137e91

Enable decorator factory pattern for `auto_axes` and `explicit_axes` API. Now users can do: `@auto_axes(...)` and `@explicit_axes(...)`

yashk2810 committed 1 year ago

440f31f6

Pallas/Mosaic support host input.

SandSnip3r committed 1 year ago

e557d4d3

Allow dropping into explicit_axes with a reshard in an outer vmap with spmd_axis_name set.

yashk2810 committed 1 year ago

f926b3cf

[Profiler] Correct xprof trace call and update tests

Matt-Hurd committed 1 year ago

72198ee2

Shard flash_attention_test

jakeharmon8 committed 1 year ago

08e67458

Remove deprecated jax.util.HashableFunction

justinjfu committed 1 year ago

4d3c4d9a

[NFC][Mosaic:TPU] Clean up elementwise rules

tlongeri committed 1 year ago

21e84c02

Fix pjit failures with direct linearize

yashk2810 committed 1 year ago

4433c140

Merge pull request #30040 from mattjj:direct-linearize-fixes-5

Google-ML-Automation committed 1 year ago

b47fa65c

Fixes `linalg.inv` + shard_map with check_vma=True by adding `standard_insert_pvary` in `linalg.solve`. Fixes https://github.com/jax-ml/jax/issues/30157

yashk2810 committed 1 year ago

a06e9e3b

[direct-linearize] fix tests to pass under JAX_USE_DIRECT_LINEARIZE=1

mattjj committed 1 year ago

e0d95928

The repr of Literal should mention the class name "Literal" in it

yashk2810 committed 1 year ago

6b495716

[pallas:mosaic] Enabled a few more primitives for all core types

superbobry committed 1 year ago

f2051a5d

[Pallas] Allow 8bit gathers

ImanHosseini committed 1 year ago

9382a794

Compute non-zero flops for `gather` with `FILL_OR_DROP`.

zacmustin committed 1 year ago

6c730874

Merge pull request #29882 from gnecula:jaxpr_consts_literalable

Google-ML-Automation committed 1 year ago

0e3fbdf5

Simplify the handling of closed-over constants in Jaxprs.

gnecula committed 1 year ago

971d34f8

[Pallas:MGPU] Fix a B200 test that was still using the recently removed direct TMEM load syntax

apaszke committed 1 year ago

b76753b8

[Mosaic GPU][NFC] Clean up pattern matching logic.

bchetioui committed 1 year ago

51437859

[Mosaic GPU][NFC] Remove duplicated code.

dimitar-asenov committed 1 year ago

0dabc520

[Pallas:MGPU] Raise a helpful error message in case someone tries to load/store to TMEM refs directly

apaszke committed 1 year ago

73e95a68

[Mosaic GPU][NFC] Do not capture op references in the Scratch manager.

dimitar-asenov committed 1 year ago

7415c23e

[Pallas:MGPU] Disallow direct reads from TMEM refs

apaszke committed 1 year ago

9c3085ba

[Mosaic:TPU][Relayouts] Materialize replicated offsets

tlongeri committed 1 year ago

7813749b

Older