ROCm/jax

Pull Requests Commits

[rollforward]: Add transfer_server_interface to xla.cc as an extra optional factory argument

pschuh committed 323 days ago

199eb3c8

Update XLA dependency to use revision

Google-ML-Automation committed 323 days ago

1ebf4d56

[Pallas/TPU] Add empty_ref_like function + SMEM support for Refs

sharadmv committed 323 days ago

fcf0b728

Rollback #30206 due to downstream test failures

Google-ML-Automation committed 323 days ago

a02bba48

Use registers for state by default in emit_pipeline

Google-ML-Automation committed 323 days ago

5efa810a

[XLA] Refactor `xla_metadata.py` to resolve dependency cycle

Google-ML-Automation committed 323 days ago

04915957

Merge pull request #30206 from carlosgmartin:logmeanexp

Google-ML-Automation committed 323 days ago

0485af20

[Mosaic] Allow sublane rotation for non-sublane dim aligned shape.

WindQAQ committed 323 days ago

11e19d1b

Reverts aee89744efbd3aa93ffea7a3bab803483add9289

jakeharmon8 committed 323 days ago

4cdd05ed

Disable `too_slow` in data.draw() for test_cast_from_32bit

jakeharmon8 committed 323 days ago

772ae661

[jax:benchmark] Add tracing benchmarks for some common operations.

danielsuo committed 323 days ago

816fab4e

Set `check_leaks=False` in `def direct_linearize` to fix some tests

yashk2810 committed 323 days ago

4752cc94

[Mosaic] Support multiple non-contracting dims if they are collapsable.

WindQAQ committed 323 days ago

0e55e131

fix jax2tf scatter impl

yashk2810 committed 323 days ago

45f67c77

[pallas] `pl.core_map` now support `functools.partial`ed functions

superbobry committed 323 days ago

841254d9

[Mosaic TPU] Fix the assumption that beforeBody's inputs size is same as afterBody's in scf.while.

bythew3i committed 323 days ago

344b1bc9

Merge pull request #30242 from Cjkkkk:support_gb300_on_cudnn_9.11+

Google-ML-Automation committed 323 days ago

cec1491d

[Mosaic GPU] Add `tmem_alloc` and `tmem_dealloc` to the Mosaic GPU dialect.

dimitar-asenov committed 324 days ago

c5790e9c

Includes `sharding_constraint_p` primitive in `roofline`.

zacmustin committed 324 days ago

af56616e

[Mosaic GPU] Add `tcgen05.mma` to the Mosaic GPU dialect.

allanrenucci committed 324 days ago

c39241ec

[Mosaic GPU] Add equational layout inference rule for `scf.WhileOp`.

bchetioui committed 324 days ago

30013df8

[Mosaic GPU] Add equational layout inference rule for `scf.ForOp`.

bchetioui committed 324 days ago

e7382abd

Update XLA dependency to use revision

Google-ML-Automation committed 324 days ago

8ff5d334

Add nn.logmeanexp.

carlosgmartin committed 324 days ago

84ab7d46

[Pallas] Physicalize fusion-dtype input avals before evaluating expressions.

Google-ML-Automation committed 324 days ago

2b510c0f

- Add more items in the ragged paged attention auto tuning table

Google-ML-Automation committed 324 days ago

c31ad6dd

add cudnn min version check on gb300(compute_cap=10.3)

Cjkkkk committed 324 days ago

b7d05023

Remove version checks for ml_dtypes >= 0.5.0

apivovarov committed 324 days ago

cdb399a1

[Pallas] Support all integer casting cases.

justinjfu committed 324 days ago

7507a3c0

[Mosaic GPU] Raise an error if TMEM allocations exceeds the limit. Currently, the kernel will hang which is difficult to debug.

justinjfu committed 324 days ago

ddc26812

Older