ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 217 days ago

9678a764

Make mosaic_gpu equation params hashable.

hawkinsp committed 218 days ago

d065d2ac

Update XLA dependency to use revision

Google-ML-Automation committed 218 days ago

e28d6ed9

[JAX] Fix the test names in colocated_python_test.py to following the standard snake case

hyeontaek committed 219 days ago

f06888f9

Remove `with_spec` from NamedSharding and replace with `.update`

yashk2810 committed 219 days ago

7904b86b

Replace `with_partitions` and `with_unreduced` with `.update` on Partitions

yashk2810 committed 219 days ago

8e346ddb

[Pallas] Add no_pipelining debugging option to emit_pipeline.

justinjfu committed 219 days ago

dfc90529

Merge pull request #29462 from hawkinsp:nvrtc

Google-ML-Automation committed 219 days ago

a849d216

Load CUDA libraries up front with cdll.LoadLibrary().

hawkinsp committed 219 days ago

080294cf

Implemented cross-host memory transfer on GPU.

mwhittaker committed 219 days ago

b9cf0af5

Make some remaining jaxpr equation params hashable.

hawkinsp committed 219 days ago

790540c7

Merge pull request #29420 from jakevdp:tpu-linalg

Google-ML-Automation committed 219 days ago

29eb8323

Merge pull request #28931 from olupton:cublas-cudnn

Google-ML-Automation committed 219 days ago

f41fc71c

Merge pull request #29456 from j-towns:pjit-p-extend-primitives

Google-ML-Automation committed 219 days ago

ec637fe4

Make params of assert_consumed_value_p hashable.

hawkinsp committed 219 days ago

64c95744

Internal refactor: move TPU lowering rules out of jax/_src/lax

jakevdp committed 219 days ago

1b1e9f71

Improve reshape not supported error message

yashk2810 committed 219 days ago

92aedb28

Make params of several pallas primitives hashable.

hawkinsp committed 219 days ago

670ae13d

[Mosaic GPU] Parametrize the `test_subview` test.

dimitar-asenov committed 219 days ago

193f11db

Fix return type annotation for tree_util.tree_broadcast.

jburnim committed 219 days ago

e66a6dd2

Add pjit_p to extend.core.primitives

j-towns committed 219 days ago

7b7a5d8c

[Mosaic GPU] Reconcile the swizzle of the a and b operands for wgmma in the Mosaic GPU dialect.

dimitar-asenov committed 219 days ago

a3542f88

Disable `too_slow` in data.draw() for test_ndindexer

jakeharmon8 committed 219 days ago

9b45aacb

[Mosaic GPU] Use warpgroup semantics for the ragged dot example kernel.

dimitar-asenov committed 219 days ago

70c90a95

[Mosaic GPU] Resolve different tile transforms using the largest common divisor.

dimitar-asenov committed 219 days ago

a4f0e402

[Mosaic GPU] Add a Mosaic GPU op `with_transforms` for manually setting memref transforms.

dimitar-asenov committed 219 days ago

fc8192ca

[Mosaic GPU] Convert all memrefs with transforms to unrealized casts and check them.

dimitar-asenov committed 219 days ago

604b6048

Move jax._src.callback to its own BUILD rule

Jake VanderPlas committed 219 days ago

142ace2c

Update XLA dependency to use revision

Google-ML-Automation committed 219 days ago

c86fefb5

fix-forward for pallas tpu memory spaces test

sharadmv committed 220 days ago

382b3e02

Older