ROCm/jax

Pull Requests Commits

Test new manylinux image

psanal35 committed 24 days ago

47083d75

[ROCm] Bump rules_ml_toolchain pin to 070c1cb (hermetic gcc-toolchain fix)

magaonka-amd committed 25 days ago

159a133d

Stop using local sysroot in ROCm build configs

psanal35 committed 25 days ago

1b5d03a1

[ROCm] Add TheRock 7.13 and TheRock latest to CI test matrices

mminutoli committed 25 days ago

28ad3828

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 25 days ago

84ed64bd

Add org-based guards to all CI caller workflows

mminutoli committed 25 days ago

b5ae6878

Merge pull request #38062 from hawkinsp:docs3

Google-ML-Automation committed 25 days ago

a5766401

Merge pull request #37977 from ROCm:use_hermetic_sysroot_for_jax_ut

Google-ML-Automation committed 25 days ago

96cd401e

Merge pull request #38010 from olupton:pbroadcast

Google-ML-Automation committed 25 days ago

54c30905

Merge pull request #37320 from samanklesaria:geometric_out_sharding

Google-ML-Automation committed 25 days ago

f941b14e

Make warnings about security on the compilation cache warnings.

hawkinsp committed 25 days ago

d5ed6d08

[pallas:mgpu] Do not use `pl.pallas_call` in :gpu_pallas_distributed_test

superbobry committed 25 days ago

f180cf4d

[pallas:sc] Do not use `mlir.lower_per_platform` for `mpmd_map`

superbobry committed 25 days ago

cf248ec2

[Mosaic GPU] Simulate commuting transforms in WG semantics.

golechwierowicz committed 25 days ago

f7fb9d27

Merge pull request #37987 from jakevdp:random-docs

Google-ML-Automation committed 25 days ago

416cb209

[NFC] Derive `shape` from `tmem_layout` in `test_load_store_tmem_scales_layout`

allanrenucci committed 25 days ago

0922eeab

Merge pull request #37254 from samanklesaria:fix_triangular

Google-ML-Automation committed 25 days ago

f9f06d42

[pallas:sc] Removed the lowering rule for `jaxpr_call`

superbobry committed 25 days ago

6fb39be2

Remove setting `Manual` HloSharding directly as it would be anyways ignored during Shardy import. If required, JAX NamedSharding should be set instead which would be correctly lowered to corresponding sdy shardings.

KanishAnand committed 25 days ago

096a1402

[Mosaic:GPU] Use multimem_parameters attribute to signal that parameter is allocated within the collective memory space.

PatriosTheGreat committed 25 days ago

17b80742

Update XLA dependency to use revision http://github.com/openxla/xla/commit/90eb26b82a4746b44d56b5aad5539679d39a33ab

Google-ML-Automation committed 25 days ago

0fd5af23

[JAX] Enable CHLO lowering for cumulative sums on GPU again.

chsigg committed 25 days ago

8f711b48

Make sure scan transposes correctly when carry is reduced.

yashk2810 committed 25 days ago

2a576d79

[Mosaic TPU] Support static lane and major dims shift of `pltpu.roll` with zero stride for packed types.

yueshengys committed 26 days ago

a1ff1d64

Error out if `*pspec` or `pspec[...]` is used when unreduced/reduced is non-empty

yashk2810 committed 26 days ago

d6c717d9

Merge pull request #38043 from jax-ml:mattjj-patch-1

Google-ML-Automation committed 26 days ago

756523e0

Fix formatting of JAX Ref documentation

mattjj committed 26 days ago

Verified cddcebdc

Add out_sharding to segment_{sum,max,min,prod}. Only allow unreduced out_sharding to `segment_sum`.

yashk2810 committed 26 days ago

ba968d94

Merge pull request #38039 from jakevdp:fix-cumsum

Google-ML-Automation committed 26 days ago

6f201d8a

Fix formatting of JAX Ref documentation

mattjj committed 26 days ago

Verified a5c1c30d

Older