ROCm/jax

Pull Requests Commits

[ROCm] Implement RNN support

Ruturaj4 committed 1 year ago

35939b96

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

d0a92c5c

Add SMEM as a supported Pallas output memory space.

Google-ML-Automation committed 1 year ago

54fd738e

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

9af29700

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

e4278f78

Merge pull request #25713 from jakevdp:debug-printoptions

Google-ML-Automation committed 1 year ago

0f4677bc

jax.debug.print: respect local np.printoptions

jakevdp committed 1 year ago

33060632

[Mosaic] NFC: Pull out vreg related functions to util.

WindQAQ committed 1 year ago

57b21541

Compute cost-analysis on only one HLO module.

zacmustin committed 1 year ago

df36c298

Merge pull request #25686 from Mikcl:docs/working-with-pytrees-formatting

Google-ML-Automation committed 1 year ago

800f903f

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

726950b8

Merge pull request #25710 from apaszke:mgpu_dialect_fix

Google-ML-Automation committed 1 year ago

68483b8e

Fix OSS build for the Mosaic GPU dialect

apaszke committed 1 year ago

64433435

[Mosaic:TPU][NFC] Clean up unused variable

tlongeri committed 1 year ago

ac817b48

Merge pull request #25706 from pearu:pearu/log10-large

Google-ML-Automation committed 1 year ago

82001ed5

Merge pull request #25661 from rdyro:tb-nightly-instructions

Google-ML-Automation committed 1 year ago

04a0fbec

Don't use x32 mode for pallas_test

apaszke committed 1 year ago

7c984c60

Reverts 83e60a9697ec20023f4e11169edf64e910b93031

apaszke committed 1 year ago

dbe9ccd6

tbp nightly instructions

rdyro committed 1 year ago

213e1782

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4a6cfebc

Fix log10 and log2 for large inputs.

pearu committed 1 year ago

50670bd9

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

e37ea587

Merge pull request #25692 from yliu120:rng_key_sharding

Google-ML-Automation committed 1 year ago

494c1573

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

25fff524

Fixes the random key sharding in shard_map.

yliu120 committed 1 year ago

97b1faac

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

879fa12d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

8eeedd18

[pallas:mosaic_gpu] Added some runtime type checking to `copy_*` and `barrier_*` primitives

superbobry committed 1 year ago

76ccb199

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7ab61b7f

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6dbda900

Older