ROCm/jax

Pull Requests Commits

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

bchetioui committed 160 days ago

542460b6

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

Google-ML-Automation committed 160 days ago

a0e1aff2

[buildcop] Skip `test_remote_profiler_gcs_path`

jakeharmon8 committed 160 days ago

0ff2bf1a

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

Google-ML-Automation committed 160 days ago

9243186b

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

superbobry committed 160 days ago

f3ea01f6

Add `remote_bytes_transferred` to Pallas `CostEstimate` dataclass.

Google-ML-Automation committed 160 days ago

e0e124b8

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

WindQAQ committed 160 days ago

f9331406

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

apaszke committed 160 days ago

f021cfa9

Fix incorrect type annotation for tensorstore.Future

Google-ML-Automation committed 160 days ago

30a27c5a

jnp.trim_zeros: support N-dimensional inputs

jakevdp committed 160 days ago

fdb4b616

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

Google-ML-Automation committed 160 days ago

7e83ca77

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

apaszke committed 160 days ago

a0f8a4aa

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

normanrink committed 160 days ago

3c7aeae3

Migrate jax2tf_test from pjit to jit

gnecula committed 160 days ago

4b5d4143

Migrate export_test to use jit instead of pjit

gnecula committed 160 days ago

3d6fa3ba

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

superbobry committed 160 days ago

3ebe0181

[Pallas:MGPU] Support reduce_scatter along arbitrary tensor dimensions

apaszke committed 160 days ago

c0165217

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

apaszke committed 160 days ago

ce52ce80

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

apaszke committed 160 days ago

32add9e6

Merge pull request #32514 from apaszke:symm2

Google-ML-Automation committed 160 days ago

300431da

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

dimitar-asenov committed 160 days ago

414f039d

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

superbobry committed 160 days ago

f0540d11

[Mosaic GPU] Fix error message.

dimitar-asenov committed 160 days ago

70d266ee

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

dimitar-asenov committed 160 days ago

260d3f0f

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

apaszke committed 160 days ago

524edf29

[Mosaic GPU] Add support for collective async copies.

allanrenucci committed 160 days ago

b4e24590

[Mosaic GPU][Serde] Add upgrade rules for vector.splat.

Google-ML-Automation committed 160 days ago

6c97ee88

Update XLA dependency to use revision http://github.com/openxla/xla/commit/48aa221e4a15067037ada8e96a58803f0601e387

Google-ML-Automation committed 160 days ago

4e14ea10

Reverts 18f96c9009624d6bec0cade1f514b8f2601f4d6a

pravnar committed 160 days ago

79231ecd

Check that exceptions are thrown when invalid types are passed to `promote_types`.

chr1sj0nes committed 160 days ago

bf57b7af

Newer Older