ROCm/jax

Pull Requests Commits

Skip ann_test test_pmap on ROCm due to IndivisibleError

AratiGanesh committed 36 days ago

59668344

[mlir] Fixed a few type errors in MLIR usages

superbobry committed 36 days ago

e67f9dc3

Change unreduced/reduced rule's signature to only take avals and params. Also combine unreduced and reduced rule because right now they have a very finicky reduced -> unreduced rule dependency.

yashk2810 committed 36 days ago

f94b79d9

[Mosaic][NFC] Rename tpu.fptosi and tpu.fptoui operand for consistency with tpu.sitofp, tpu.uitofp and the arith dialect ops

tlongeri committed 36 days ago

9ce09a0a

Merge pull request #35533 from magaonka-amd:fix/skip-aot-test-on-gpu

Google-ML-Automation committed 36 days ago

1dbd1ea2

Merge pull request #34722 from ROCm:tot-rocm-skips

Google-ML-Automation committed 36 days ago

10957831

PR #35842: Updated Pyrefly to 0.56

superbobry committed 36 days ago

8ad2e838

internal change

Google-ML-Automation committed 36 days ago

ee6208fa

Make PyLoadedExecutable's destructor call `SetDeleteOptions(thread_id)`.

Google-ML-Automation committed 36 days ago

babbce66

[pallas:mosaic] Moved `get_device_kind` and `get_num_device_cores` into `tpu_info`

superbobry committed 36 days ago

94fcb020

Merge pull request #35779 from jakevdp:pyrefly-jax2tf

Google-ML-Automation committed 36 days ago

4852eee8

[Mosaic TPU] Fix some tests behavior.

yueshengys committed 36 days ago

132bfbac

Merge pull request #35823 from superbobry:mlir

Google-ML-Automation committed 37 days ago

5d46c524

[jaxlib] Added type stubs for a few more MLIR-related extensions

superbobry committed 37 days ago

572082bb

[pyrefly] add coverage for jax2tf, excluding tests

jakevdp committed 37 days ago

99ab770f

Merge pull request #35787 from jakevdp:pyrefly-experimental-2

Google-ML-Automation committed 37 days ago

6466a62a

Merge pull request #35745 from jax-ml:no-lu-shard-map

Google-ML-Automation committed 37 days ago

70b0a6e3

Update rules_ml_toolchain to integrate new CUPTI headers.

ybaturina committed 37 days ago

12553810

[pyrefly] type-check jax.experimental/array_serialization & colocated_python

jakevdp committed 37 days ago

180b30f3

Add constraints on output during cond, while_loop, sort lowering

yashk2810 committed 37 days ago

877ca781

shard map without linear util

dougalm committed 37 days ago

e58a022a

Merge pull request #35775 from jakevdp:pyrefly-experimental

Google-ML-Automation committed 37 days ago

0936349f

[Mosaic TPU] Strengthen the tpu.memref_slice verifier

apaszke committed 37 days ago

fc635497

Increase size of `ode_test`.

mwhittaker committed 37 days ago

ef13248c

[pallas:mosaic] Fixed the signature of `infer_tiling` -- it never returns `None`

superbobry committed 37 days ago

bcb3e200

[NFC] Extract reduce expression reduction logic into a helper function.

allanrenucci committed 37 days ago

fc2e5146

[Mosaic GPU] Remove workaround for miscompilation

apaszke committed 37 days ago

c1f4c33a

[Mosaic GPU] Add support for atomic stores from arrays with the WG strided layout

apaszke committed 37 days ago

42c4d0a3

Update XLA dependency to use revision http://github.com/openxla/xla/commit/0e27d63f1902324b1cd8cae442d14a29aa36c0f3

Google-ML-Automation committed 37 days ago

4b578528

[Mosaic GPU] Clarify semantic of `BroadcastInDimOp`.

allanrenucci committed 37 days ago

336ca0c8

Older