ROCm/jax

Pull Requests Commits

Merge pull request #32586 from mattjj:hijax-compiled

Google-ML-Automation committed 137 days ago

a9b9cf3c

Add a fallback to `get_hlo_text()` to go via old code if AttributeError is hit.

Google-ML-Automation committed 137 days ago

1aa217b8

[buildcop] Fix acc_smem in matmul1

jakeharmon8 committed 137 days ago

5f351a0b

Fix typo in _one_hot

justinjfu committed 137 days ago

41bb2d7e

[Mosaic GPU] Bump the forward compatible version of Mosaic GPU to 2.

bchetioui committed 137 days ago

542460b6

[Pallas/interpreter] Make `races` (of type `RaceDetectionState`) a global object of the interpreter (again).

Google-ML-Automation committed 137 days ago

a0e1aff2

[buildcop] Skip `test_remote_profiler_gcs_path`

jakeharmon8 committed 137 days ago

0ff2bf1a

Merge pull request #32594 from jakevdp:trim-zeros-nd

Google-ML-Automation committed 137 days ago

9243186b

[pallas:sc] Added support for `pl.multiple_of`

superbobry committed 137 days ago

f3ea01f6

Add `remote_bytes_transferred` to Pallas `CostEstimate` dataclass.

Google-ML-Automation committed 137 days ago

e0e124b8

[Mosaic] Force relayout on int vregs where pack/unpack support is more comprehensive.

WindQAQ committed 137 days ago

f9331406

[Pallas:MGPU] Only infer vec_size=1 in the reduce scatter kernel for integer types

apaszke committed 137 days ago

f021cfa9

Fix incorrect type annotation for tensorstore.Future

Google-ML-Automation committed 137 days ago

30a27c5a

jnp.trim_zeros: support N-dimensional inputs

jakevdp committed 137 days ago

fdb4b616

Merge pull request #32614 from normanrink:mosaic-interpret-fixes

Google-ML-Automation committed 137 days ago

7e83ca77

[Pallas:MGPU] Add an all-reduce mode to the reduce-scatter kernel

apaszke committed 137 days ago

a0f8a4aa

[Pallas/interpreter] Add dependency on `__init__.py` source file to target.

normanrink committed 137 days ago

3c7aeae3

Migrate jax2tf_test from pjit to jit

gnecula committed 137 days ago

4b5d4143

Migrate export_test to use jit instead of pjit

gnecula committed 137 days ago

3d6fa3ba

[pallas:mosaic] Fixed the `NDIndexer` instantiation in `_prng_key_load_lowering_rule`

superbobry committed 137 days ago

3ebe0181

[Pallas:MGPU] Support reduce_scatter along arbitrary tensor dimensions

apaszke committed 137 days ago

c0165217

[Pallas:MGPU] Try to collapse dimension pairs in transfer_strided

apaszke committed 137 days ago

ce52ce80

[Pallas:MGPU] Add a small benchmark script to the reduce scatter kernel

apaszke committed 137 days ago

32add9e6

Merge pull request #32514 from apaszke:symm2

Google-ML-Automation committed 138 days ago

300431da

[Mosaic GPU] Add an SMEM inference rule for `memref.transpose`.

dimitar-asenov committed 138 days ago

414f039d

Use `()` instead of `tuple()` and `{}` instead of `dict()`

superbobry committed 138 days ago

f0540d11

[Mosaic GPU] Fix error message.

dimitar-asenov committed 138 days ago

70d266ee

[Mosaic GPU] Add SMEM inference rule for `mgpu.custom_op`.

dimitar-asenov committed 138 days ago

260d3f0f

[Mosaic GPU] Add a basic test for running collectives on Torch tensors

apaszke committed 138 days ago

524edf29

[Mosaic GPU] Add support for collective async copies.

allanrenucci committed 138 days ago

b4e24590

Newer Older