ROCm/jax

Pull Requests Commits

[Pallas:MGPU] Expose multicast TMA stores

apaszke committed 98 days ago

69fd2052

[Pallas:MGPU] Add an all-gather implementation

apaszke committed 98 days ago

8eed7dc5

Add tpu_sc.subcore_barrier primitive.

brianwa84 committed 98 days ago

6aae6f43

Add support for -1 size to ref reshaper, for convenience.

brianwa84 committed 98 days ago

8674d771

[Pallas:MGPU] Add ReducedLayout support to inline_mgpu params

Google-ML-Automation committed 98 days ago

5608cb7c

[pallas:sc] Skip `pl.multiple_of` test on old libtpu

superbobry committed 98 days ago

fff4d58d

[pallas:sc] Renamed the `pl.debug_check` test to match other SC tests

superbobry committed 98 days ago

8ade11bf

Update XLA dependency to use revision http://github.com/openxla/xla/commit/39fbe6fa509bd4c93fc35e7a9690b6e283fe634b

Google-ML-Automation committed 98 days ago

174025a6

Reverts 4ee9cb2bdf89fdd6e05b43b77e8cb593464e0a56

Google-ML-Automation committed 98 days ago

fc0c1d08

Switch ShapeDtypeStruct to use concrete_mesh when it's created with PartitionSpec during construction time.

yashk2810 committed 98 days ago

20e84d0a

make a "hello world" of pallas_call + pinned buffers

mattjj committed 98 days ago

50e2a612

Merge pull request #32586 from mattjj:hijax-compiled

Google-ML-Automation committed 98 days ago

a9b9cf3c

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

Google-ML-Automation committed 98 days ago

1aa217b8

[buildcop] Fix acc_smem in matmul1

jakeharmon8 committed 98 days ago

5f351a0b

Fix typo in _one_hot

justinjfu committed 98 days ago

41bb2d7e

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

bchetioui committed 98 days ago

542460b6

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

Google-ML-Automation committed 98 days ago

a0e1aff2

[buildcop] Skip `test_remote_profiler_gcs_path`

jakeharmon8 committed 98 days ago

0ff2bf1a

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

Google-ML-Automation committed 98 days ago

9243186b

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

superbobry committed 98 days ago

f3ea01f6

Add `remote_bytes_transferred` to Pallas `CostEstimate` dataclass.

Google-ML-Automation committed 98 days ago

e0e124b8

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

WindQAQ committed 98 days ago

f9331406

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

apaszke committed 98 days ago

f021cfa9

Fix incorrect type annotation for tensorstore.Future

Google-ML-Automation committed 98 days ago

30a27c5a

jnp.trim_zeros: support N-dimensional inputs

jakevdp committed 98 days ago

fdb4b616

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

Google-ML-Automation committed 98 days ago

7e83ca77

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

apaszke committed 98 days ago

a0f8a4aa

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

normanrink committed 98 days ago

3c7aeae3

Migrate jax2tf_test from pjit to jit

gnecula committed 99 days ago

4b5d4143

Migrate export_test to use jit instead of pjit

gnecula committed 99 days ago

3d6fa3ba

Newer Older