ROCm/jax

Pull Requests Commits

relax version requirements for ROCm Jax Plugin wheels, allow overriding using env var

JehandadKhan committed 89 days ago

d5c7e430

Add a explicit_x64_dtypes flag.

hawkinsp committed 89 days ago

9b6df1dc

[Pallas:MGPU] Expose cluster launch controls in Pallas, allowing an SM to dynamically query the next unit of work to schedule.

justinjfu committed 89 days ago

1c817be6

Integrate LLVM at llvm/llvm-project@79a0bf0efce5

Google-ML-Automation committed 89 days ago

262b80c9

[Pallas][Mosaic GPU] Use NDLoopInfo as the argument to the loop body of plgpu.nd_loop

justinjfu committed 89 days ago

b6996cbd

[pallas:sc] Added overloads for `parallel_loop` and slightly reformatted the docstring

superbobry committed 89 days ago

85edead0

[Mosaic GPU] Rename `transfer_tiled2` to `transfer_tiled` and remove previously unused `transfer_tiled` method.

bchetioui committed 89 days ago

c531c419

[Mosaic GPU] Infer transforms for load/store from/to registers with MMA layout.

allanrenucci committed 89 days ago

a5b388f2

[Mosaic GPU][NFC] Consolidate `is_memref_transposed` helper functions.

allanrenucci committed 89 days ago

4792072c

Adds `fuser.custom_fusion` for defining ops with custom `fuser.fuse` behavior.

jburnim committed 89 days ago

5c916c2b

Don't serialize MLIR modules unnecessarily when passing them to the Python bindings.

hawkinsp committed 89 days ago

5b1ddeff

[Pallas][interpreter] Add an interpreter option for allowing allocations of (shared) HBM buffers under `run_scoped`.

Google-ML-Automation committed 89 days ago

ef026e49

[state] Allowed refs as indexers

superbobry committed 89 days ago

7d5a09a5

Merge pull request #32095 from superbobry:main

Google-ML-Automation committed 89 days ago

3adb4ad7

Bumped mypy to 1.18.2 and ruff to 0.13.1

superbobry committed 89 days ago

ac6ea955

Update XLA dependency to use revision http://github.com/openxla/xla/commit/43f2d5b6dc6bb381f189013c63c585934e0fb60c

Google-ML-Automation committed 89 days ago

f2d90b08

Merge pull request #32085 from mattjj:vmap-primitive

Google-ML-Automation committed 90 days ago

de355b17

Include `_pathways.so` in jaxlib.

lukebaumann committed 90 days ago

20dcbcfb

[hijax] add a flag gating whether vmap is a primitive

mattjj committed 90 days ago

edfcff5b

[hijax] impelment a vmap_p primitive, and its consequences

mattjj committed 90 days ago

af54a78e

Mark args with jax.buffer_donor if output with same size exists

ICGog committed 90 days ago

4d8b7010

Merge pull request #30133 from carlosgmartin:tree_reduce_unspecified

Google-ML-Automation committed 90 days ago

c86b7514

Merge pull request #32076 from ROCm:jd/wsl_experimental_rebased

Google-ML-Automation committed 90 days ago

360eacc5

Update documentation to reflect experimental support for WSL2 on ROCm

JehandadKhan committed 90 days ago

3629d460

[Pallas:SC] Improve error message for pl.squeezed

brianwa84 committed 90 days ago

366a6789

Internal change

yashk2810 committed 90 days ago

266be973

Fix the 7x mesh order for 2x2x2 topology

yashk2810 committed 90 days ago

4997988b

Fix the number of SparseCores per generation

naummo committed 90 days ago

e67d5005

Reland GPU compute host offloading tests.

basioli-k committed 90 days ago

d107ddc1

[Mosaic GPU][NFC] Add `SMEMTiling` layout.

dimitar-asenov committed 90 days ago

4864c71c

Older