ROCm/jax

Pull Requests Commits

[Pallas Fuser] Add basic sublane/lane reshape fusion

sharadmv committed 237 days ago

6951cb92

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

sharadmv committed 237 days ago

0f169a2e

sharadmv committed 237 days 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 237 days ago

3b3c1b3e

Add out_sharding on broadcast_to().

pschuh committed 237 days ago

4cd63a7c

Merge pull request #28467 from mattjj:list-get

Google-ML-Automation committed 237 days ago

f05edc7f

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

superbobry committed 237 days ago

f7252497

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

mattjj committed 237 days ago

11d33824

remove spurious assertion

mattjj committed 237 days ago

d985f366

Merge pull request #28464 from jax-ml:shmap_doc

Google-ML-Automation committed 237 days ago

7d395c20

[state] Slightly restructured `_is_trivial_indexer`

superbobry committed 237 days ago

23df262e

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

Google-ML-Automation committed 237 days ago

cbf79e0a

Update shard_map.md's API specification section

yashk2810 committed 237 days ago

42572338

Integrate LLVM at llvm/llvm-project@7752e0a10b25

Google-ML-Automation committed 237 days ago

5a81db00

Merge pull request #28318 from mattjj:boxes3

Google-ML-Automation committed 237 days ago

e7d65f99

Wait for async result to populate the exception.

Google-ML-Automation committed 237 days 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 237 days ago

b59c0e48

Propagate function name when recording elapsed time event.

danielsuo committed 237 days ago

879eb41c

Merge pull request #28408 from Cjkkkk:fix_return_residual_train

Google-ML-Automation committed 237 days ago

18b0bc64

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

Google-ML-Automation committed 237 days ago

d226cf77

Fix bug in all_leaves when is_leaf is specified

Luke Tsekouras committed 237 days ago

4768e003

jax._src.util: improve type annotations

jakevdp committed 238 days 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 238 days ago

48001a24

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

superbobry committed 238 days ago

dcdc25ba

Rename with_dll_constraint to with_layout_constraint for the upcoming Layout API rename!

yashk2810 committed 238 days ago

9107c634

[Mosaic GPU] Add a debug_print for TMEM refs

apaszke committed 238 days ago

e189cd46

Fix ASAN/MSAN/TSAN failures for buffer callback.

dfm committed 238 days ago

37b87b52

Update XLA dependency to use revision

Google-ML-Automation committed 238 days ago

c0da649d

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

apaszke committed 238 days ago

0e1c6a4b

Skip buffer_callback tests on TPU.

dfm committed 238 days ago

c47bc83e

Older