ROCm/jax

Pull Requests Commits

Point nightly ROCm wheel tests at pinned TheRock 10.0.0rc2/rc4 images.

psanal35 committed 46 days ago

be0e3642

Merge pull request #40219 from mattjj:push-skysvottzsln

Google-ML-Automation committed 46 days ago

151feb1a

[docs] move host-offloading content to 201/memory-spaces.md, other moves

mattjj committed 46 days ago

d739b7a5

Merge pull request #40217 from mattjj:push-lrxvwsozvkuo

Google-ML-Automation committed 46 days ago

5e291607

Merge pull request #40216 from mattjj:push-zqsmnmuwooyr

Google-ML-Automation committed 46 days ago

8c191b84

[docs] move up pallas link in sidebar, tweaks

mattjj committed 46 days ago

238cc90c

[docs] move convolution tutorial into 101 docs

mattjj committed 46 days ago

38b246a5

Merge pull request #40213 from mattjj:push-sltsvksyznkx

Google-ML-Automation committed 46 days ago

ca64b720

[docs] move all pytree content to 101

mattjj committed 46 days ago

6d492b57

Merge pull request #40212 from mattjj:push-qsuyyktyvqmp

Google-ML-Automation committed 46 days ago

918445ee

[mosaic] Don't emit an equation in _advance_slot when the predicate is statically known.

hawkinsp committed 46 days ago

714f015c

[docs] move jax-internals docs into their own 601 section

mattjj committed 46 days ago

cc75acb5

[mosaic] In the pipeline emitter, check for and short circuit identical objects in _tuples_differ.

hawkinsp committed 46 days ago

b5af1f3b

Merge pull request #40211 from mattjj:push-qnlmnwmxtsws

Google-ML-Automation committed 46 days ago

1ea3227a

[docs] add a "what's new" docs page to highlight things we haven't before

mattjj committed 46 days ago

9cbc36be

Merge pull request #40186 from Intel-tensorflow:kranipa/oneapi_hybrid_kernels

Google-ML-Automation committed 46 days ago

f4f89eab

Deduplicate `GetIfrtHloSharding`

junwhanahn committed 46 days ago

a9e471e2

Don't allow `Tracer.memory_space` to work. Please use `jax.typeof(x).memory_space`.

Cristian Garcia committed 46 days ago

abd868d6

Merge pull request #40203 from mattjj:push-quxutlrnsklo

Google-ML-Automation committed 46 days ago

747d59bb

Add OneAPI Hybrid Kernels

kranipa committed 46 days ago

4d68253f

Make the leveled docs (101-501) the primary documentation

mattjj committed 46 days ago

26007fdd

Merge pull request #40182 from mattjj:push-lsmklnzoomxo

Google-ML-Automation committed 46 days ago

e76122f9

[Pallas/GPU interpreter] Relax tolerances.

bchetioui committed 46 days ago

2a38bfc6

Merge pull request #40193 from hawkinsp:reqs

Google-ML-Automation committed 46 days ago

a097066d

Remove Python 3.15 version constraints on requirements files.

hawkinsp committed 46 days ago

67759b66

Roll back cuDNN upgrade.

bchetioui committed 46 days ago

5d131506

[Pallas/Mosaic GPU] Initial support for layouts in types in Pallas/Mosaic GPU.

bchetioui committed 46 days ago

29ca0c0e

Adjust `test_fun_sourceinfo` to be compatible with Windows.

bchetioui committed 46 days ago

b472879a

Use float16 for input array in memories_test.py.

gnecula committed 46 days ago

e83c4525

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

Google-ML-Automation committed 46 days ago

d1a8f5c9

Older