ROCm/jax

Pull Requests Commits

update wheel_version_suffix custom

kiran-thumma committed 40 days ago

Verified b8eb358f

tag name jaxlib whl

kiran-thumma committed 40 days ago

Verified 2dcfe7a4

Reverts f44bc6e37f770a9d79d24061f7960235a2892407

jakeharmon8 committed 75 days ago

403977d2

Prepare for JAX release 0.8.0

jakeharmon8 committed 75 days ago

b1ce4c57

[Pallas] Fix 64-bit multiplication emulation in Philox kernel

justinjfu committed 75 days ago

07dad8e8

Update XLA hash for JAX release

jakeharmon8 committed 75 days ago

4f06e7a6

Fix race when dialects were registered in an MLIR context when it was being accessed concurrently.

hawkinsp committed 75 days ago

1bd3402c

Clean up dead code.

brianwa84 committed 75 days ago

b942e894

Remove cusolver_potrf_ffi from the list of covered backwards compatibility targets.

gnecula committed 75 days ago

ac87cbe0

[Pallas:MGPU] Expose multicast TMA stores

apaszke committed 75 days ago

69fd2052

[Pallas:MGPU] Add an all-gather implementation

apaszke committed 75 days ago

8eed7dc5

Add tpu_sc.subcore_barrier primitive.

brianwa84 committed 75 days ago

6aae6f43

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

brianwa84 committed 75 days ago

8674d771

[Pallas:MGPU] Add ReducedLayout support to inline_mgpu params

Google-ML-Automation committed 75 days ago

5608cb7c

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

superbobry committed 76 days ago

fff4d58d

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

superbobry committed 76 days ago

8ade11bf

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

Google-ML-Automation committed 76 days ago

174025a6

Reverts 4ee9cb2bdf89fdd6e05b43b77e8cb593464e0a56

Google-ML-Automation committed 76 days ago

fc0c1d08

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

yashk2810 committed 76 days ago

20e84d0a

make a "hello world" of pallas_call + pinned buffers

mattjj committed 76 days ago

50e2a612

Merge pull request #32586 from mattjj:hijax-compiled

Google-ML-Automation committed 76 days ago

a9b9cf3c

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

Google-ML-Automation committed 76 days ago

1aa217b8

[buildcop] Fix acc_smem in matmul1

jakeharmon8 committed 76 days ago

5f351a0b

Fix typo in _one_hot

justinjfu committed 76 days ago

41bb2d7e

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

bchetioui committed 76 days ago

542460b6

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

Google-ML-Automation committed 76 days ago

a0e1aff2

[buildcop] Skip `test_remote_profiler_gcs_path`

jakeharmon8 committed 76 days ago

0ff2bf1a

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

Google-ML-Automation committed 76 days ago

9243186b

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

superbobry committed 76 days ago

f3ea01f6

Add `remote_bytes_transferred` to Pallas `CostEstimate` dataclass.

Google-ML-Automation committed 76 days ago

e0e124b8

Older