ROCm/jax

Pull Requests Commits

Switch ROCm RBE to linux_x64_gpu_do_gfx950 pool

charleshofer committed 47 days ago

5379d003

Merge pull request #40223 from mattjj:push-ryoqplqkztyu

Google-ML-Automation committed 47 days ago

9e2b382f

[typing] Use PEP 695 typevars in jaxlib.pytree

Jake VanderPlas committed 47 days ago

425dfea7

[docs] highlight the new (hijax-based) ffi docs

mattjj committed 47 days ago

206840cb

Merge pull request #40219 from mattjj:push-skysvottzsln

Google-ML-Automation committed 48 days ago

151feb1a

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

mattjj committed 48 days ago

d739b7a5

Merge pull request #40217 from mattjj:push-lrxvwsozvkuo

Google-ML-Automation committed 48 days ago

5e291607

Merge pull request #40216 from mattjj:push-zqsmnmuwooyr

Google-ML-Automation committed 48 days ago

8c191b84

[docs] move up pallas link in sidebar, tweaks

mattjj committed 48 days ago

238cc90c

[docs] move convolution tutorial into 101 docs

mattjj committed 48 days ago

38b246a5

Merge pull request #40213 from mattjj:push-sltsvksyznkx

Google-ML-Automation committed 48 days ago

ca64b720

[docs] move all pytree content to 101

mattjj committed 48 days ago

6d492b57

Merge pull request #40212 from mattjj:push-qsuyyktyvqmp

Google-ML-Automation committed 48 days ago

918445ee

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

hawkinsp committed 48 days ago

714f015c

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

mattjj committed 48 days ago

cc75acb5

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

hawkinsp committed 48 days ago

b5af1f3b

Merge pull request #40211 from mattjj:push-qnlmnwmxtsws

Google-ML-Automation committed 48 days ago

1ea3227a

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

mattjj committed 48 days ago

9cbc36be

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

Google-ML-Automation committed 48 days ago

f4f89eab

Deduplicate `GetIfrtHloSharding`

junwhanahn committed 48 days ago

a9e471e2

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

Cristian Garcia committed 48 days ago

abd868d6

Merge pull request #40203 from mattjj:push-quxutlrnsklo

Google-ML-Automation committed 48 days ago

747d59bb

Add OneAPI Hybrid Kernels

kranipa committed 48 days ago

4d68253f

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

mattjj committed 48 days ago

26007fdd

Merge pull request #40182 from mattjj:push-lsmklnzoomxo

Google-ML-Automation committed 48 days ago

e76122f9

[Pallas/GPU interpreter] Relax tolerances.

bchetioui committed 48 days ago

2a38bfc6

Merge pull request #40193 from hawkinsp:reqs

Google-ML-Automation committed 48 days ago

a097066d

Remove Python 3.15 version constraints on requirements files.

hawkinsp committed 48 days ago

67759b66

Roll back cuDNN upgrade.

bchetioui committed 48 days ago

5d131506

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

bchetioui committed 48 days ago

29ca0c0e

Older