ROCm/jax

Pull Requests Commits

[Mosaic GPU] Execute relayouts in the lowering of `layout_cast`.

bchetioui committed 346 days ago

a6f03d3d

Major deps update:

vam-google committed 346 days ago

d7671692

Removed remaining reference to `TritonCompilerParams`

superbobry committed 346 days ago

dbcb8407

Add traceme to `PythonRefManager::CollectGarbage`

junwhanahn committed 346 days ago

71ea45bb

Add some tracemes in py_array to make slow device put debugging easier.

Google-ML-Automation committed 346 days ago

d46202e4

[mutable-arrays] re-land https://github.com/jax-ml/jax/pull/29353

mattjj committed 346 days ago

bb4fbc54

[mosaic] `MemRef{Slice,Squeeze}` verifiers now support strided layouts

superbobry committed 346 days ago

f3370cb5

[pallas:mosaic] Fixed a typo in the distributed tutorial

superbobry committed 346 days ago

346ce85d

[Pallas/Mosaic GPU] Propagate transforms on the accumulator in `tcgen05.mma`.

bchetioui committed 346 days ago

3fdc97b8

Remove `Layout`, `.layout`, `.input_layouts` and `.output_layouts` and replace it with `Format`, `.format`, `.input_formats` and `.output_formats` in JAX

yashk2810 committed 346 days ago

bfc07e2f

[Mosaic GPU][NFC] Add `checkInLayouts` and `checkOutLayouts` utils to `gpu_layout_inference_test.py`.

bchetioui committed 347 days ago

e818940d

[Mosaic GPU] Change layout inference tests to rely on explicit `layout_cast`s.

bchetioui committed 347 days ago

84066b77

[Pallas:MGPU] Add docs for pl.core_map and plgpu.kernel

apaszke committed 347 days ago

f99d2b49

Create a test suite for Pallas mosaic GPU tests.

dimitar-asenov committed 347 days ago

b99d004c

[Mosaic GPU] Delete dead code in `layout_inference.py`.

bchetioui committed 347 days ago

e55f55fc

Reenable AVX512 after LLVM fix upstream.

WillFroom committed 347 days ago

0b54a1e8

Pass shardy option through jax config.

ZixuanJiang committed 347 days ago

1e4a0f76

Merge pull request #28872 from Cjkkkk:jax_cudnn_sdpa_mla

Google-ML-Automation committed 347 days ago

4efa56d2

Merge pull request #28135 from justinjfu:gpu_pipe_docs

Google-ML-Automation committed 347 days ago

4d2b14a4

Merge pull request #29434 from jakevdp:normalize-indices

Google-ML-Automation committed 347 days ago

00852217

add cudnn sdpa mla support

Cjkkkk committed 347 days ago

d1883667

[JAX] Skip failing tpu tests until June 30th.

subhankarshah committed 347 days ago

9d1b01e0

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 347 days ago

a47ae573

Merge pull request #29560 from kanglant:update_lock_files

Google-ML-Automation committed 347 days ago

d3f08713

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 347 days ago

5e290ddd

disable cudnn sdpa on blackwell gb300

Cjkkkk committed 347 days ago

0a5731da

Merge pull request #29563 from jakevdp:fix-doc

Google-ML-Automation committed 347 days ago

8337fe5b

[doc] fix build error

jakevdp committed 347 days ago

6567192a

Merge pull request #29554 from dubstack:patch-1

Google-ML-Automation committed 347 days ago

8cbd915b

Update jax requirements lock files after 0.6.2 release

kanglant committed 347 days ago

cb2315a8

Newer Older