ROCm/jax

Pull Requests Commits

Update documentation to reflect experimental support for WSL2 on ROCm

JehandadKhan committed 1 year ago

3629d460

Internal change

yashk2810 committed 1 year ago

266be973

Fix the 7x mesh order for 2x2x2 topology

yashk2810 committed 1 year ago

4997988b

Fix the number of SparseCores per generation

naummo committed 1 year ago

e67d5005

Reland GPU compute host offloading tests.

basioli-k committed 1 year ago

d107ddc1

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

dimitar-asenov committed 1 year ago

4864c71c

Enable GPU mosaic B200 test targets

IllogicalMoose committed 1 year ago

c096089d

[pallas:sc] Simplified the `lax.bitcast_convert_type` test

superbobry committed 1 year ago

a02d6a74

[Mosaic GPU][NFC] Remove duplicate import.

allanrenucci committed 1 year ago

ea8f44e3

[Pallas:MGPU] Use a much better matmul kernel in the collective matmul

apaszke committed 1 year ago

8f21ab30

[Pallas:SC] Add new tpu_sc primitive for cumulative sum.

brianwa84 committed 1 year ago

40836799

[Mosaic:TPU] Have VectorLayout::join handle fully replicated layouts and fix bug in infer-vector-layout elementwise rule

tlongeri committed 1 year ago

02cad83c

[pallas:sc] Do not emit trivial bitcasts to the same dtype

superbobry committed 1 year ago

62b5112d

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

Google-ML-Automation committed 1 year ago

f2dc38d2

Adds support for carried scalar and vector values in plsc.parallel_loop

brianwa84 committed 1 year ago

7cc283a5

Reverts 7dc02949020edf52cdae1aa8e2707a109cf4511d

Google-ML-Automation committed 1 year ago

8309eae7

Revert eager constant folding in `lax._convert_element_type`.

hawkinsp committed 1 year ago

6a3f1e3c

Follow-up after [RBE Docker image update](https://github.com/openxla/xla/pull/31443).

Google-ML-Automation committed 1 year ago

20f328c5

Reverts 94356fb726cd6ea21bdbd4ba13414ac003c4d288

Google-ML-Automation committed 1 year ago

fae2101f

[Mosaic] Canonicalize scalar integers ops

Google-ML-Automation committed 1 year ago

69b1729c

Flag for bounds checking changed.

Google-ML-Automation committed 1 year ago

f3393106

Remove JAX_GPU_HAVE_SPARSE.

hawkinsp committed 1 year ago

6c1760b6

[pallas:sc] Added support for int32 to `plsc.{pack,unpack}`

superbobry committed 1 year ago

9f8e65a8

Integrate Triton up to

vwbaker committed 1 year ago

18268f0b

Fix remaining tests under `pmap_shmap_merge=True`.

danielsuo committed 1 year ago

d8ac2546

Adding split_by_mesh_axis to experimental for use by a new experimental reshard.

lukebaumann committed 1 year ago

9de08133

[Pallas/Mosaic GPU][NFC] Clarify `wait_smem_to_gmem` docstring.

bchetioui committed 1 year ago

f67b62f5

Merge pull request #31994 from mattbahr:implement-hyp2f1-derivative

Google-ML-Automation committed 1 year ago

c3aa8692

[Mosaic] Refactor: Move op-specific rules first

Google-ML-Automation committed 1 year ago

94356fb7

Lift the restriction where lowering with only AbstractMesh required all inputs to have the same AbstractMesh. Now we only require the size of AbstractMesh to match across all inputs.

yashk2810 committed 1 year ago

f5288192

Older