ROCm/jax

Pull Requests Commits

[MosaicGPU] Extract code into a new method `BarrierRef.from_dialect_barrier_memref` and implement support for 1D barrier memrefs.

dimitar-asenov committed 1 year ago

22417ae2

[Mosaic GPU] Allow querying layouts from a `FuncOp`'s block arguments if set.

bchetioui committed 1 year ago

bc7204f0

[mosaic] Extracted serialization pass traversal logic into a reusable function

superbobry committed 1 year ago

4221f109

Merge pull request #25917 from ROCm:ci_fix_multi_gpu_test_logic-upstream

Google-ML-Automation committed 1 year ago

9a60e6fc

[Pallas TPU] Add helpers file with copy_ref function

sharadmv committed 1 year ago

0ac63157

[Mosaic] Allow passing `ApplyVectorLayoutCtx` to tpu.apply_layout_op.

WindQAQ committed 1 year ago

4a9cc9ff

Fix run_multi_gpu script multi-gpu issue and refactor code

Ruturaj4 committed 1 year ago

8e88adcd

[Mosaic] Add a macro to convert abseil StatusOr to LLVM FailureOr

naummo committed 1 year ago

d3ba1eb3

Move halt for testing step to be just before running tests

nitins17 committed 1 year ago

8a053af1

Merge pull request #25906 from ROCm:ci_add_new_gfx-upstream

Google-ML-Automation committed 1 year ago

cf67e28f

Merge pull request #25911 from hawkinsp:version

Google-ML-Automation committed 1 year ago

2fa10020

Make utils for reporting function name work with `functools.partial` by using the inner `.func` attribute if the object doesn't have a `__name__` attribute. `functools.partial` objects do not have `__name__` attributes by default.

ZacharyGarrett committed 1 year ago

f7d097f7

Update the JAX version to 0.5.0.

hawkinsp committed 1 year ago

3a8f31aa

Merge pull request #25755 from ROCm:ci_rnn_final-upstream

Google-ML-Automation committed 1 year ago

41993fdb

Merge pull request #25864 from jax-ml:yet-more-linearization-fixes

Google-ML-Automation committed 1 year ago

ca012d7a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

51f23100

Jax: Stop returning a list of cost-analyses.

zacmustin committed 1 year ago

2d72e8de

Merge pull request #25876 from gnecula:debug_info_3

Google-ML-Automation committed 1 year ago

70c1ee5d

Add gfx12xx archs

Ruturaj4 committed 1 year ago

435edf1f

Merge pull request #25880 from jakevdp:fix-gather

Google-ML-Automation committed 1 year ago

2e5e4799

More linearization fixes

dougalm committed 1 year ago

9fe553ca

[pallas:mosaic_gpu] Fixed a crash in MLIR Python bindings

superbobry committed 1 year ago

afcb21dd

[Mosaic GPU][NFC] Address some previous stylistic comments.

bchetioui committed 1 year ago

cdf490a5

[Pallas TPU] Temporarily strengthen restrictions on Pallas tests

apaszke committed 1 year ago

aa19f9c4

[better_errors] More cleanup

gnecula committed 1 year ago

f9dfe7f6

[pallas] Fix bad rebase, deleted lowering for a print

Google-ML-Automation committed 1 year ago

c4406d27

[pallas][mosaic kernel export] Add initial support for exporting a dynamic shapes (placeholder bound) kernel out of mosaic, via pallas as both MLIR and jaxpr.

Google-ML-Automation committed 1 year ago

c18492be

[ROCm] Implement RNN support

Ruturaj4 committed 1 year ago

fe68eb8b

[Pallas] Fix GQA triton kernel test.

justinjfu committed 1 year ago

cc9f6e75

Temporarily disable GQA attention tests on GPU, which were broken by a Triton integrate.

hawkinsp committed 1 year ago

d1810b42

Newer Older