ROCm/jax

Pull Requests Commits

Raise an error when unreduced/reduced is used with the GSPMD partitioner.

yashk2810 committed 210 days ago

55e77803

[JAX] Fix flakiness in colocated_python_test

hyeontaek committed 210 days ago

e0d9b0bf

[Mosaic] Constant vector factory for brevity

naummo committed 210 days ago

ca8eebfc

[Mosaic] i32 test helpers for brevity

naummo committed 210 days ago

ca71572f

Remove info logging that's not necessary.

yashk2810 committed 210 days ago

b30d7e51

[Mosaic] Set TensorCore as the default kernel core

naummo committed 210 days ago

ebe34614

Add `select_n` to `roofline`.

zacmustin committed 210 days ago

d2d090b8

[JAX] Preserve `DeviceList` during a call to `batched_copy_array_to_devices_with_sharding()`

junwhanahn committed 210 days ago

b996073c

Update XLA dependency to use revision

danielsuo committed 210 days ago

e6b2ab7e

Move jax._src.error_check to its own build rule

Jake VanderPlas committed 210 days ago

8213f607

Prevent negative index behavior by using unsigned index.

Google-ML-Automation committed 210 days ago

210f151b

[Mosaic GPU] Use TiledLayout for TMEM layouts

apaszke committed 211 days ago

2c382f86

[jax:pallas] Update fusible matmul test tolerances.

danielsuo committed 211 days ago

c4ba59bf

Remove deprecated LayoutUtil::IsDenseArray

apivovarov committed 211 days ago

9901084e

Move `XLA_COMMIT` and `XLA_SHA256` Bazel symbols to their own file.

danielsuo committed 211 days ago

b8b8c308

Fix lax.switch where unmapped arg was not being broadcasted on the correct `axis_data.explicit_mesh_axis` under vmap.

yashk2810 committed 211 days ago

f7faaa81

Reverts 4f871b0e2cbf0148580d68589173facb02a9e1ac

MichaelHudgins committed 211 days ago

30ba5c9c

Merge pull request #29649 from mattjj:attribute-error-fixes

Google-ML-Automation committed 211 days ago

1acf8ff7

[Mosaic TPU] Clean up transpose canonicalization

apaszke committed 211 days ago

d33f74f2

Update XLA dependency to use revision

hawkinsp committed 211 days ago

9d158b0f

don't pass `self` to `raise AttributeError(...)`

mattjj committed 211 days ago

bbd41ae0

Automated Code Change

Google-ML-Automation committed 211 days ago

cf60242a

Remove `DeviceLocalLayout` and `.device_local_layout` aliases.

yashk2810 committed 211 days ago

c057c7cd

Move jax._src.shard_map to its own BUILD rule

Jake VanderPlas committed 211 days ago

353363ae

Move jax._src.flatten_util to its own build rule

Jake VanderPlas committed 211 days ago

5749a34d

Finish `jax._src.lax` BUILD refactor

Jake VanderPlas committed 211 days ago

526ff073

Merge pull request #29626 from jakevdp:jax-numpy-imports

Google-ML-Automation committed 211 days ago

726a4094

Make canonicalization of reshapes more robust:

Google-ML-Automation committed 211 days ago

7d924e8f

Check for existence of all `None`s in a spec to consider a sharding as replicated in reshape.

yashk2810 committed 212 days ago

ab8720c2

Use `Layout` in place of `DeviceLocalLayout`

yashk2810 committed 212 days ago

88b0747a

Older