ROCm/jax

Pull Requests Commits

Run the TPU workflow on new self-hosted runners

nitins17 committed 1 year ago

7d7a0fa2

[shape_poly] Adding shape polymorphism support for the state primitives.

gnecula committed 1 year ago

0831e2e3

Integrate LLVM at llvm/llvm-project@33fcd6acc755

metaflow committed 1 year ago

1bc9df42

[pallas:mosaic_gpu] Pulled `delay_release` into `emit_pipeline`

superbobry committed 1 year ago

f18df8f3

Reverts c04aec9d525dd2e767495e41b98e82dd79315f37

naummo committed 1 year ago

e72b4490

[sharding_in_types] Add `concatenate_p` support

yashk2810 committed 1 year ago

6568713a

[Mosaic TPU] Add bound check for general vector store op.

bythew3i committed 1 year ago

869a5334

[sharding_in_types] Add `pad_p` support to sharding_in_types to handle transpose to slice correctly.

yashk2810 committed 1 year ago

840cf3f7

[Pallas] Use Pallas cost estimator for flash attention.

justinjfu committed 1 year ago

1f6152d1

Merge pull request #25020 from jakevdp:lax-pad-validation

Google-ML-Automation committed 1 year ago

f39392ea

Adds Google Sans font

barnesjoseph committed 1 year ago

bf7f9aa8

jax.lax.pad: improve input validation

jakevdp committed 1 year ago

17825882

Merge pull request #25019 from jakevdp:lax-pad-doc

Google-ML-Automation committed 1 year ago

334bd4d0

DOC: add examples for jax.lax.pad

jakevdp committed 1 year ago

2699e950

[array api] use most recent version of array_api_tests

jakevdp committed 1 year ago

f749fca7

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6fe78042

[sharding_in_types] Add slice_p and squeeze_p sharding rule to make flash attention work in backward pass

yashk2810 committed 1 year ago

9b941808

Merge pull request #25011 from jakevdp:jnp-linalg-module

Google-ML-Automation committed 1 year ago

d219439d

[Pallas TPU] Support masked store

bythew3i committed 1 year ago

9d2f62f8

[sharding_in_types] Make flash_attention forward pass in TPU pallas work nicely with sharding in types. Backward pass is still busted which I will fix in follow up CLs.

yashk2810 committed 1 year ago

40fc6598

Merge pull request #25013 from hawkinsp:relnotes

Google-ML-Automation committed 1 year ago

19b4996e

Mention stackless in the release notes.

hawkinsp committed 1 year ago

dfe27a16

Merge pull request #25008 from skye:barrier

Google-ML-Automation committed 1 year ago

1a3e693a

Set __module__ attribute of jax.numpy.linalg APIs

jakevdp committed 1 year ago

621e39de

[pallas:mosaic_gpu] Avoid using multiple indexers in the parallel grid test

superbobry committed 1 year ago

9584ee3b

Don't psum over auto mesh dims in _unmentioned2.

pschuh committed 1 year ago

2c9b917b

Remove internal KeyArray alias

jakevdp committed 1 year ago

fee272e5

Merge pull request #25004 from jax-ml:linearize-trace

Google-ML-Automation committed 1 year ago

eab9026c

Make a direct linearize trace.

dougalm committed 1 year ago

d0f17c0c

[pallas mgpu] Lowering for while loops as long as they are secretly for loops.

cperivol committed 1 year ago

8d84f283

Newer Older