ROCm/jax

Pull Requests Commits

[Pallas Fuser] Add basic sublane/lane reshape fusion

sharadmv committed 1 year ago

6951cb92

[Pallas/Fuser] Ignore ops that have no_block_spec being pulled

sharadmv committed 1 year ago

0f169a2e

sharadmv committed 1 year ago

6e580231

Add out_sharding to `jnp.ravel` and `jnp.reshape` and `jnp.dot` which just forward out_sharding to their lax variants.

yashk2810 committed 1 year ago

3b3c1b3e

Add out_sharding on broadcast_to().

pschuh committed 1 year ago

4cd63a7c

Merge pull request #28467 from mattjj:list-get

Google-ML-Automation committed 1 year ago

f05edc7f

[mosaic_gpu] Added support for using `cf.assert` in Mosaic GPU kernels

superbobry committed 1 year ago

f7252497

[attrs] replace list.freeze() with list.get(), only allow locals

mattjj committed 1 year ago

11d33824

remove spurious assertion

mattjj committed 1 year ago

d985f366

Merge pull request #28464 from jax-ml:shmap_doc

Google-ML-Automation committed 1 year ago

7d395c20

[state] Slightly restructured `_is_trivial_indexer`

superbobry committed 1 year ago

23df262e

Enable usage of mirrored `.tar` redistributions in Bazel RBE GPU jobs.

Google-ML-Automation committed 1 year ago

cbf79e0a

Update shard_map.md's API specification section

yashk2810 committed 1 year ago

42572338

Integrate LLVM at llvm/llvm-project@7752e0a10b25

Google-ML-Automation committed 1 year ago

5a81db00

Merge pull request #28318 from mattjj:boxes3

Google-ML-Automation committed 1 year ago

e7d65f99

Wait for async result to populate the exception.

Google-ML-Automation committed 1 year ago

02c2768d

Initial commit for `smap` i.e. shard_map1D. The signature is `smap(f, in_axes, out_axes, axis_name)`. This change does NOT make the API public.

yashk2810 committed 1 year ago

b59c0e48

Propagate function name when recording elapsed time event.

danielsuo committed 1 year ago

879eb41c

Merge pull request #28408 from Cjkkkk:fix_return_residual_train

Google-ML-Automation committed 1 year ago

18b0bc64

Merge pull request #28413 from jakevdp:jax-util-annotations

Google-ML-Automation committed 1 year ago

d226cf77

Fix bug in all_leaves when is_leaf is specified

Luke Tsekouras committed 1 year ago

4768e003

jax._src.util: improve type annotations

jakevdp committed 1 year ago

42d76fb3

Make literal's dtype print with an empty shape so that it's consistent. So `1.0:f32` -> `1.0:f32[]`

yashk2810 committed 1 year ago

48001a24

[mosaic_gpu] Use `jtu` helpers instead of `get_sass`

superbobry committed 1 year ago

dcdc25ba

Rename with_dll_constraint to with_layout_constraint for the upcoming Layout API rename!

yashk2810 committed 1 year ago

9107c634

[Mosaic GPU] Add a debug_print for TMEM refs

apaszke committed 1 year ago

e189cd46

Fix ASAN/MSAN/TSAN failures for buffer callback.

dfm committed 1 year ago

37b87b52

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c0da649d

[Mosaic GPU] Add support for passing in tcgen05.mma LHS in TMEM

apaszke committed 1 year ago

0e1c6a4b

Skip buffer_callback tests on TPU.

dfm committed 1 year ago

c47bc83e

Older