ROCm/jax

Pull Requests Commits

[ROCm] Implement RNN support

Ruturaj4 committed 352 days ago

35939b96

Update XLA dependency to use revision

Google-ML-Automation committed 352 days ago

d0a92c5c

Add SMEM as a supported Pallas output memory space.

Google-ML-Automation committed 353 days ago

54fd738e

Update XLA dependency to use revision

Google-ML-Automation committed 353 days ago

9af29700

Update XLA dependency to use revision

Google-ML-Automation committed 354 days ago

e4278f78

Merge pull request #25713 from jakevdp:debug-printoptions

Google-ML-Automation committed 355 days ago

0f4677bc

jax.debug.print: respect local np.printoptions

jakevdp committed 355 days ago

33060632

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

WindQAQ committed 355 days ago

57b21541

Compute cost-analysis on only one HLO module.

zacmustin committed 355 days ago

df36c298

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

Google-ML-Automation committed 355 days ago

800f903f

Update XLA dependency to use revision

Google-ML-Automation committed 355 days ago

726950b8

Merge pull request #25710 from apaszke:mgpu_dialect_fix

Google-ML-Automation committed 355 days ago

68483b8e

Fix OSS build for the Mosaic GPU dialect

apaszke committed 355 days ago

64433435

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

tlongeri committed 355 days ago

ac817b48

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

Google-ML-Automation committed 355 days ago

82001ed5

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

Google-ML-Automation committed 355 days ago

04a0fbec

Don't use x32 mode for pallas_test

apaszke committed 355 days ago

7c984c60

Reverts 83e60a9697ec20023f4e11169edf64e910b93031

apaszke committed 355 days ago

dbe9ccd6

tbp nightly instructions

rdyro committed 356 days ago

213e1782

Update XLA dependency to use revision

Google-ML-Automation committed 356 days ago

4a6cfebc

Fix log10 and log2 for large inputs.

pearu committed 357 days ago

50670bd9

Update XLA dependency to use revision

Google-ML-Automation committed 357 days ago

e37ea587

Merge pull request #25692 from yliu120:rng_key_sharding

Google-ML-Automation committed 358 days ago

494c1573

Update XLA dependency to use revision

Google-ML-Automation committed 358 days ago

25fff524

Fixes the random key sharding in shard_map.

yliu120 committed 359 days ago

97b1faac

Update XLA dependency to use revision

Google-ML-Automation committed 359 days ago

879fa12d

Update XLA dependency to use revision

Google-ML-Automation committed 360 days ago

8eeedd18

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

superbobry committed 360 days ago

76ccb199

Update XLA dependency to use revision

Google-ML-Automation committed 361 days ago

7ab61b7f

Update XLA dependency to use revision

Google-ML-Automation committed 362 days ago

6dbda900

Older