ROCm/jax

Pull Requests Commits

Add traceme to `PythonRefManager::CollectGarbage`

junwhanahn committed 1 year ago

71ea45bb

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

Google-ML-Automation committed 1 year ago

d46202e4

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

superbobry committed 1 year ago

f3370cb5

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

superbobry committed 1 year ago

346ce85d

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

bchetioui committed 1 year 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 1 year ago

bfc07e2f

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

bchetioui committed 1 year ago

e818940d

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

bchetioui committed 1 year ago

84066b77

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

apaszke committed 1 year ago

f99d2b49

Create a test suite for Pallas mosaic GPU tests.

dimitar-asenov committed 1 year ago

b99d004c

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

bchetioui committed 1 year ago

e55f55fc

Reenable AVX512 after LLVM fix upstream.

WillFroom committed 1 year ago

0b54a1e8

Pass shardy option through jax config.

ZixuanJiang committed 1 year ago

1e4a0f76

Merge pull request #28872 from Cjkkkk:jax_cudnn_sdpa_mla

Google-ML-Automation committed 1 year ago

4efa56d2

Merge pull request #28135 from justinjfu:gpu_pipe_docs

Google-ML-Automation committed 1 year ago

4d2b14a4

Merge pull request #29434 from jakevdp:normalize-indices

Google-ML-Automation committed 1 year ago

00852217

add cudnn sdpa mla support

Cjkkkk committed 1 year ago

d1883667

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

subhankarshah committed 1 year ago

9d1b01e0

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 1 year ago

a47ae573

Merge pull request #29560 from kanglant:update_lock_files

Google-ML-Automation committed 1 year ago

d3f08713

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 1 year ago

5e290ddd

Merge pull request #29563 from jakevdp:fix-doc

Google-ML-Automation committed 1 year ago

8337fe5b

[doc] fix build error

jakevdp committed 1 year ago

6567192a

Merge pull request #29554 from dubstack:patch-1

Google-ML-Automation committed 1 year ago

8cbd915b

Update jax requirements lock files after 0.6.2 release

kanglant committed 1 year ago

cb2315a8

Merge pull request #29550 from hawkinsp:py311

Google-ML-Automation committed 1 year ago

09121436

Fix bugs in the double_buffered_pipeline example

dubstack committed 1 year ago

4c54c022

Move jax._src.lax to its own BUILD rule

Jake VanderPlas committed 1 year ago

511bf2fe

Fix rare error with Literal in DynamicJaxprTracer.full_lower.

jburnim committed 1 year ago

7c9613a6

Skip pytest for Python 3.14 during the JAX release process

kanglant committed 1 year ago

4f871b0e

Older