ROCm/jax

Pull Requests Commits

Add traceme to `PythonRefManager::CollectGarbage`

junwhanahn committed 353 days ago

71ea45bb

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

Google-ML-Automation committed 353 days ago

d46202e4

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

superbobry committed 353 days ago

f3370cb5

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

superbobry committed 354 days ago

346ce85d

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

bchetioui committed 354 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 354 days ago

bfc07e2f

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

bchetioui committed 354 days ago

e818940d

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

bchetioui committed 354 days ago

84066b77

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

apaszke committed 354 days ago

f99d2b49

Create a test suite for Pallas mosaic GPU tests.

dimitar-asenov committed 354 days ago

b99d004c

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

bchetioui committed 354 days ago

e55f55fc

Reenable AVX512 after LLVM fix upstream.

WillFroom committed 354 days ago

0b54a1e8

Pass shardy option through jax config.

ZixuanJiang committed 354 days ago

1e4a0f76

Merge pull request #28872 from Cjkkkk:jax_cudnn_sdpa_mla

Google-ML-Automation committed 354 days ago

4efa56d2

Merge pull request #28135 from justinjfu:gpu_pipe_docs

Google-ML-Automation committed 354 days ago

4d2b14a4

Merge pull request #29434 from jakevdp:normalize-indices

Google-ML-Automation committed 354 days ago

00852217

add cudnn sdpa mla support

Cjkkkk committed 354 days ago

d1883667

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

subhankarshah committed 354 days ago

9d1b01e0

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 354 days ago

a47ae573

Merge pull request #29560 from kanglant:update_lock_files

Google-ML-Automation committed 354 days ago

d3f08713

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 354 days ago

5e290ddd

Merge pull request #29563 from jakevdp:fix-doc

Google-ML-Automation committed 354 days ago

8337fe5b

[doc] fix build error

jakevdp committed 354 days ago

6567192a

Merge pull request #29554 from dubstack:patch-1

Google-ML-Automation committed 354 days ago

8cbd915b

Update jax requirements lock files after 0.6.2 release

kanglant committed 354 days ago

cb2315a8

Merge pull request #29550 from hawkinsp:py311

Google-ML-Automation committed 354 days ago

09121436

Fix bugs in the double_buffered_pipeline example

dubstack committed 354 days ago

4c54c022

Move jax._src.lax to its own BUILD rule

Jake VanderPlas committed 354 days ago

511bf2fe

Fix rare error with Literal in DynamicJaxprTracer.full_lower.

jburnim committed 354 days ago

7c9613a6

Skip pytest for Python 3.14 during the JAX release process

kanglant committed 354 days ago

4f871b0e

Older