ROCm/jax

Pull Requests Commits

[ROCm] Add TheRock containers to CI test matrices

mminutoli committed 23 days ago

585ae354

Fix `so_lib` path with NVSHMEM lib.

ybaturina committed 23 days ago

c0622a5d

Merge pull request #36658 from eamartin:always_throw_for_inconsistent_cache

Google-ML-Automation committed 23 days ago

eee721ad

create SingleDeviceSharding when possible

Google-ML-Automation committed 23 days ago

cdcfe9d8

Always raise in case of CacheVerificationError

eamartin committed 23 days ago

cdafedaa

[Mosaic] Add xla::Array broadcast util function

tlongeri committed 23 days ago

9ac02bef

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

yashk2810 committed 23 days ago

6c98690f

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

yashk2810 committed 23 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 23 days ago

b1843782

[Pallas] Add core_map tests to tpu_fusible_matmul_test

jburnim committed 23 days ago

03aa6f43

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

Google-ML-Automation committed 23 days ago

7cb36e57

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

Cristian Garcia committed 23 days ago

caaf49c8

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

sharadmv committed 23 days ago

8756a172

Update hermetic Python references to @rules_ml_toolchain

yuriivcs committed 23 days ago

97269056

[Mosaic] Canonicalize tpu.dynamic_gather of singleton dimension

tlongeri committed 23 days ago

7203f8e8

[doc] fix irregularities in `jax.scipy.stats` docstrings

Jake VanderPlas committed 23 days ago

0b8afd08

[consts] Optimize the sharding of the hoisted constant arguments

gnecula committed 24 days ago

c4918e04

Fix flaky deprecation warning in pjit_test.

mwhittaker committed 24 days ago

af134295

Fix the residuals returned by xla_metadata_call_linearize rule

yashk2810 committed 24 days ago

16972660

`is_fully_replicated` should return False if the underlying array is unreduced on any mesh axes.

yashk2810 committed 24 days ago

9b677fc1

Merge pull request #36974 from gnecula:consts_min_size

Google-ML-Automation committed 24 days ago

cafc4576

[Mosaic:GPU] Replace jtu.test_device_matches(["cuda"]) with jtu.is_device_cuda().

PatriosTheGreat committed 24 days ago

3919f0bd

[pallas:sc] Removed leftover `debug=True` in tests

superbobry committed 24 days ago

d9e9f83e

Merge pull request #36959 from gnecula:jax2tf_pyrefly

Google-ML-Automation committed 24 days ago

5ac1cad8

Removed leftover jaxlib version guards

superbobry committed 24 days ago

2f36e9df

Add a new stage_p primitive and Trace.stage_value operation.

hawkinsp committed 24 days ago

92faff13

[jax2tf] Add pyrefly suppressions

gnecula committed 24 days ago

83c1eff2

Merge pull request #36936 from jakevdp:key-bool

Google-ML-Automation committed 24 days ago

de15bbeb

[Pallas SC] Support indexers that have both indirect offsets and minor slices

apaszke committed 24 days ago

5112029a

[consts] Allow small constants to be included in the HLO and executable.

gnecula committed 24 days ago

4468ff41

Older