ROCm/jax

Pull Requests Commits

[ROCm][Pallas] Skip fused attention bwd tests that exceed device resources

mminutoli committed 104 days ago

d7fa5ce2

Trim down release validation jobs for testing.

tsrw2048 committed 105 days ago

4e4bce83

Add expanded set of specific arch targets to support.

tsrw2048 committed 105 days ago

80c90840

Point JAX to downstream XLA and bump XLA commit

tsrw2048 committed 105 days ago

5d50bd7e

Add HSA_NO_SCRATCH_RECLAIM env var and update Rock images to 7.14 in release workflow

tsrw2048 committed 105 days ago

5c870c42

Added a release validation workflow for JAX ROCm wheels.

tsrw2048 committed 105 days ago

776f8a90

Cherry pick env var leak fix from open upstream.

tsrw2048 committed 105 days ago

c1876694

Remove outdated ROCm skip that depends on ROCm versions < 6.4 (#776)

gulsumgudukbay committed 105 days ago

3d72c19b

[ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images

gulsumgudukbay committed 105 days ago

d525d049

[ROCm] Set HSA_NO_SCRATCH_RECLAIM=1 in ROCm test workflows

gulsumgudukbay committed 105 days ago

20f77158

[ROCm] Embed TheRock-aware wheel RUNPATH at link time

gulsumgudukbay committed 105 days ago

6bf8ce5e

Add org-based guards to all CI caller workflows

mminutoli committed 105 days ago

2bf7b3d7

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 105 days ago

cf269722

[ROCm] Add `remote_download_outputs` option at build invocation.

tsrw2048 committed 105 days ago

f3fd9a7b

[ROCm] xdist: select GPU from runner-allocated ROCR slice (#794)

mminutoli committed 109 days ago

Verified 42886f9b

[ROCm] Use grid=(1,) in pallas test_non_range_while_loop

magaonka-amd committed 131 days ago

cbcc4e01

Update rules_ml_toolchain reference

alekstheod committed 145 days ago

b2e32655

Override HIP_VISIBLE_DEVICES per ROCm xdist worker

magaonka-amd committed 145 days ago

45b7af75

[ROCm] Split ROCm pytest into single- and multi-accelerator passes

magaonka-amd committed 145 days ago

95aba756

Use CloudFront download for ROCm plugin/PJRT wheels instead of S3

psanal35 committed 145 days ago

acd00f93

Cherry-pick: Fix visibility in BUILD files (#757) (#759)

magaonka-amd committed 145 days ago

a23d3768

Add a libtpu version guard in test_dce_sink_prevents_xla_dce

yashk2810 committed 160 days ago

a33ed614

Remove unused EXTRA_INDEX environment variable

nitins17 committed 160 days ago

79e8ef93

Allow shell globbing for libtpu wheel path in pytest_tpu.yml

nitins17 committed 160 days ago

83e0ac50

Prepare for JAX release 0.10.0

yashk2810 committed 160 days ago

08a228a9

Fix a bug in WeakValueInterner.

hawkinsp committed 160 days ago

2362cc76

Addressed an old pytype TODO

superbobry committed 160 days ago

c12e417e

Skip a test in OSS that relies on stdout

Rifur13 committed 160 days ago

1d0523d6

[Pallas] Make VMEM explicit on pallas_call TPU kernel inputs during tracing.

jburnim committed 160 days ago

01c0f46f

Use @immutable for ShapeDtypeStruct and object.__setattr__ in the constructor.

Google-ML-Automation committed 160 days ago

ad793d6a

Older