ROCm/jax

Pull Requests Commits

Don't eagerly compute error messages, especially in assertAllClose.

hawkinsp committed 76 days ago

5ff3a9e8

[Mosaic GPU] Add layout inference rule to AsymcStoreSparseMetadataSmemToTmem op.

golechwierowicz committed 76 days ago

00c40ca8

Apply the documentation fix from https://github.com/jax-ml/jax/pull/20828 using math notation.

hawkinsp committed 76 days ago

f06655b9

[Mosaic GPU] Add AsyncStoreSparseMetadataSmemToTmemOp op.

golechwierowicz committed 76 days ago

69996b16

Apply a fix from PR https://github.com/jax-ml/jax/pull/21020.

hawkinsp committed 76 days ago

24f387da

Merge pull request #35855 from jakevdp:array-overrides

Google-ML-Automation committed 76 days ago

7303b4e1

Remove jax/tools/colab_tpu.py.

hawkinsp committed 76 days ago

67da5f94

Fix an inappropriate use of /tmp in a test on Windows.

hawkinsp committed 76 days ago

9d8d5bc5

Merge pull request #35856 from superbobry:pyrefly-2

Google-ML-Automation committed 76 days ago

41d3f9fd

Merge pull request #35841 from gnecula:jax2tf_named_sharding

Google-ML-Automation committed 76 days ago

726c2244

[mlir] Fixed a few type errors in MLIR usages

superbobry committed 76 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 76 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 76 days ago

9ce09a0a

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

Google-ML-Automation committed 76 days ago

1dbd1ea2

Removed unnecessary jaxlib version guard

superbobry committed 76 days ago

9d471312

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

Google-ML-Automation committed 76 days ago

10957831

[maint] clean up array class overrides

jakevdp committed 76 days ago

b2466104

PR #35842: Updated Pyrefly to 0.56

superbobry committed 76 days ago

8ad2e838

internal change

Google-ML-Automation committed 76 days ago

ee6208fa

Address pr comments

gulsumgudukbay committed 76 days ago

0a602824

[jax2tf] Move to use the exported NamedShardings

gnecula committed 76 days ago

f8f4bbd1

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

Google-ML-Automation committed 77 days ago

babbce66

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

superbobry committed 77 days ago

94fcb020

Merge pull request #35779 from jakevdp:pyrefly-jax2tf

Google-ML-Automation committed 77 days ago

4852eee8

[Mosaic TPU] Fix some tests behavior.

yueshengys committed 77 days ago

132bfbac

Merge pull request #35823 from superbobry:mlir

Google-ML-Automation committed 77 days ago

5d46c524

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

superbobry committed 77 days ago

572082bb

[pyrefly] add coverage for jax2tf, excluding tests

jakevdp committed 77 days ago

99ab770f

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

Google-ML-Automation committed 77 days ago

6466a62a

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

Google-ML-Automation committed 77 days ago

70b0a6e3

Newer Older