ROCm/jax

Pull Requests Commits

Update documentation to reflect experimental support for WSL2 on ROCm

JehandadKhan committed 122 days ago

3629d460

Internal change

yashk2810 committed 122 days ago

266be973

Fix the 7x mesh order for 2x2x2 topology

yashk2810 committed 122 days ago

4997988b

Fix the number of SparseCores per generation

naummo committed 122 days ago

e67d5005

Reland GPU compute host offloading tests.

basioli-k committed 122 days ago

d107ddc1

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

dimitar-asenov committed 122 days ago

4864c71c

Enable GPU mosaic B200 test targets

IllogicalMoose committed 122 days ago

c096089d

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

superbobry committed 122 days ago

a02d6a74

[Mosaic GPU][NFC] Remove duplicate import.

allanrenucci committed 122 days ago

ea8f44e3

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

apaszke committed 122 days ago

8f21ab30

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

brianwa84 committed 122 days ago

40836799

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

tlongeri committed 122 days ago

02cad83c

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

superbobry committed 122 days ago

62b5112d

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

Google-ML-Automation committed 122 days ago

f2dc38d2

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

brianwa84 committed 122 days ago

7cc283a5

Reverts 7dc02949020edf52cdae1aa8e2707a109cf4511d

Google-ML-Automation committed 122 days ago

8309eae7

Revert eager constant folding in `lax._convert_element_type`.

hawkinsp committed 122 days ago

6a3f1e3c

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

Google-ML-Automation committed 122 days ago

20f328c5

Reverts 94356fb726cd6ea21bdbd4ba13414ac003c4d288

Google-ML-Automation committed 122 days ago

fae2101f

[Mosaic] Canonicalize scalar integers ops

Google-ML-Automation committed 122 days ago

69b1729c

Flag for bounds checking changed.

Google-ML-Automation committed 122 days ago

f3393106

Remove JAX_GPU_HAVE_SPARSE.

hawkinsp committed 122 days ago

6c1760b6

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

superbobry committed 123 days ago

9f8e65a8

Integrate Triton up to

vwbaker committed 123 days ago

18268f0b

Fix remaining tests under `pmap_shmap_merge=True`.

danielsuo committed 123 days ago

d8ac2546

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

lukebaumann committed 123 days ago

9de08133

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

bchetioui committed 123 days ago

f67b62f5

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

Google-ML-Automation committed 123 days ago

c3aa8692

[Mosaic] Refactor: Move op-specific rules first

Google-ML-Automation committed 123 days 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 123 days ago

f5288192

Older