ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

9678a764

Make mosaic_gpu equation params hashable.

hawkinsp committed 1 year ago

d065d2ac

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

e28d6ed9

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

hyeontaek committed 1 year ago

f06888f9

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

yashk2810 committed 1 year ago

7904b86b

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

yashk2810 committed 1 year ago

8e346ddb

[Pallas] Add no_pipelining debugging option to emit_pipeline.

justinjfu committed 1 year ago

dfc90529

Merge pull request #29462 from hawkinsp:nvrtc

Google-ML-Automation committed 1 year ago

a849d216

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

hawkinsp committed 1 year ago

080294cf

Implemented cross-host memory transfer on GPU.

mwhittaker committed 1 year ago

b9cf0af5

Make some remaining jaxpr equation params hashable.

hawkinsp committed 1 year ago

790540c7

Merge pull request #29420 from jakevdp:tpu-linalg

Google-ML-Automation committed 1 year ago

29eb8323

Merge pull request #28931 from olupton:cublas-cudnn

Google-ML-Automation committed 1 year ago

f41fc71c

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

Google-ML-Automation committed 1 year ago

ec637fe4

Make params of assert_consumed_value_p hashable.

hawkinsp committed 1 year ago

64c95744

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

jakevdp committed 1 year ago

1b1e9f71

Improve reshape not supported error message

yashk2810 committed 1 year ago

92aedb28

Make params of several pallas primitives hashable.

hawkinsp committed 1 year ago

670ae13d

[Mosaic GPU] Parametrize the `test_subview` test.

dimitar-asenov committed 1 year ago

193f11db

Fix return type annotation for tree_util.tree_broadcast.

jburnim committed 1 year ago

e66a6dd2

Add pjit_p to extend.core.primitives

j-towns committed 1 year ago

7b7a5d8c

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

dimitar-asenov committed 1 year ago

a3542f88

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

jakeharmon8 committed 1 year ago

9b45aacb

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

dimitar-asenov committed 1 year ago

70c90a95

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

dimitar-asenov committed 1 year ago

a4f0e402

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

dimitar-asenov committed 1 year ago

fc8192ca

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

dimitar-asenov committed 1 year ago

604b6048

Move jax._src.callback to its own BUILD rule

Jake VanderPlas committed 1 year ago

142ace2c

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c86fefb5

fix-forward for pallas tpu memory spaces test

sharadmv committed 1 year ago

382b3e02

Older