ROCm/jax

Pull Requests Commits

add jax[rocm-local] extra and relax plugin version check for accepting wildcards

gulsumgudukbay committed 20 hours ago

2fb2d6cf

internal change

Google-ML-Automation committed 1 day ago

ee6208fa

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

Google-ML-Automation committed 2 days ago

babbce66

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

superbobry committed 2 days ago

94fcb020

Merge pull request #35779 from jakevdp:pyrefly-jax2tf

Google-ML-Automation committed 2 days ago

4852eee8

[Mosaic TPU] Fix some tests behavior.

yueshengys committed 2 days ago

132bfbac

Merge pull request #35823 from superbobry:mlir

Google-ML-Automation committed 2 days ago

5d46c524

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

superbobry committed 2 days ago

572082bb

[pyrefly] add coverage for jax2tf, excluding tests

jakevdp committed 2 days ago

99ab770f

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

Google-ML-Automation committed 2 days ago

6466a62a

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

Google-ML-Automation committed 2 days ago

70b0a6e3

Update rules_ml_toolchain to integrate new CUPTI headers.

ybaturina committed 2 days ago

12553810

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

jakevdp committed 2 days ago

180b30f3

Add constraints on output during cond, while_loop, sort lowering

yashk2810 committed 2 days ago

877ca781

shard map without linear util

dougalm committed 2 days ago

e58a022a

Merge pull request #35775 from jakevdp:pyrefly-experimental

Google-ML-Automation committed 2 days ago

0936349f

[Mosaic TPU] Strengthen the tpu.memref_slice verifier

apaszke committed 2 days ago

fc635497

Increase size of `ode_test`.

mwhittaker committed 2 days ago

ef13248c

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

superbobry committed 2 days ago

bcb3e200

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

allanrenucci committed 2 days ago

fc2e5146

[Mosaic GPU] Remove workaround for miscompilation

apaszke committed 2 days ago

c1f4c33a

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

apaszke committed 2 days ago

42c4d0a3

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

Google-ML-Automation committed 2 days ago

4b578528

[Mosaic GPU] Clarify semantic of `BroadcastInDimOp`.

allanrenucci committed 2 days ago

336ca0c8

[Pallas:MGPU] Add a test for `broadcast_in_dim` of mixed size-1 and size-0 dimensions.

allanrenucci committed 2 days ago

3071c638

[NFC] Remove redundant `num_threads` and `thread_name` parameters in tests.

allanrenucci committed 2 days ago

2baa52ba

[Mosaic TPU] Specialize indices of grid dims of size 1 to static 0

apaszke committed 2 days ago

ac71172e

[Pallas:TPU] Drop the buffer count to 1 when the pltpu.emit_pipeline operand window is constant

apaszke committed 3 days ago

1d2ef4cc

[Mosaic GPU] Add support for keeping dimensions in `Reduce` constraint.

allanrenucci committed 3 days ago

53346b39

[Mosaic GPU] Add bitwidth to IsTransferable.

golechwierowicz committed 3 days ago

b94ce497

Older