ROCm/jax

Pull Requests Commits

Test ROCm 7.2.2 update for build-artifacts for continuous

psanal35 committed 175 days ago

fc7b9a72

Test ROCm 7.2.2 update

psanal35 committed 175 days ago

71aebc75

Upgrade upstream ROCm CI from 7.2.0 to 7.2.2

Ruturaj4 committed 175 days ago

0a48872f

[Mosaic GPU] Add `shape` and `memory_space` properties to `constraints.Variable`.

allanrenucci committed 181 days ago

05a69ef5

[pallas:sc] Removed leftover debug print

superbobry committed 181 days ago

c2b61060

Merge pull request #36975 from gnecula:log_consts

Google-ML-Automation committed 181 days ago

82f4a4a7

[Pallas] Update changelog after release of JAX 0.10.0.

bchetioui committed 181 days ago

5f5c1367

Update XLA dependency to use revision http://github.com/openxla/xla/commit/07015e0a17e028acb8da7f74ecb8d3d51d35f807

Google-ML-Automation committed 181 days ago

f97e0648

Automated Code Change

Google-ML-Automation committed 182 days ago

588aac94

Merge `modify_sdy_sharding_wrt_axis_types` into `_to_sdy_sharding`

yashk2810 committed 182 days ago

347e9702

Merge pull request #37012 from jakevdp:is-scalar-util

Google-ML-Automation committed 182 days ago

1bf6ffa5

Make SdyArrayList frozen and cache _to_sdy_sharding

yashk2810 committed 182 days ago

e644eefe

Call get_owner half as many times during MLIR lowering.

hawkinsp committed 182 days ago

3d2a1f98

[Pallas TPU] Change CoreMemorySpace to pair a memory space and a mesh

sharadmv committed 182 days ago

a6565a8c

[internal] rename is_python_scalar

jakevdp committed 182 days ago

43c22960

Make SdyArray immutable and cache SdyArray to sharding attribute conversion.

hawkinsp committed 182 days ago

21c3eb04

Simplify the flattening of TransformedRefs in DMAs.

IvyZX committed 182 days ago

b6fb583e

Expose timeouts to python.

pschuh committed 182 days ago

bba277dd

Fix `so_lib` path with NVSHMEM lib.

ybaturina committed 182 days ago

c0622a5d

Merge pull request #36658 from eamartin:always_throw_for_inconsistent_cache

Google-ML-Automation committed 182 days ago

eee721ad

create SingleDeviceSharding when possible

Google-ML-Automation committed 182 days ago

cdcfe9d8

Always raise in case of CacheVerificationError

eamartin committed 182 days ago

cdafedaa

[Mosaic] Add xla::Array broadcast util function

tlongeri committed 182 days ago

9ac02bef

De-duplicate skipping 1-sized mesh axis code in tracing/lowering.

yashk2810 committed 182 days ago

6c98690f

Remove `sharding_proto` from `lower_with_sharding_in_types` since it's not needed anymore

yashk2810 committed 182 days ago

2ae423ac

Add a private parameter `_manual_axis_type` on `lax.empty` so that pl.kernel/core_map can propagate `manual_axis_type` on refs correctly.

yashk2810 committed 182 days ago

b1843782

[Pallas] Add core_map tests to tpu_fusible_matmul_test

jburnim committed 182 days ago

03aa6f43

Merge pull request #36851 from magaonka-amd:feature/rocm-tpu-style-pytest

Google-ML-Automation committed 182 days ago

7cb36e57

Propagate requires_low and is_high in process_custom_{jvp,vjp}_call

Cristian Garcia committed 182 days ago

caaf49c8

[Pallas] Rename out_shapes/scratch_shapes to out_types/scratch_types in mpmd_map

sharadmv committed 182 days ago

8756a172

Older