ROCm/jax

Pull Requests Commits

[ROCm] Added a release validation workflow for JAX ROCm wheels.

tsrw2048 committed 12 hours ago

cae86fbe

Add org-based guards to all CI caller workflows

mminutoli committed 13 hours ago

5e905aa6

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 13 hours ago

8a61f1fb

Add all specific target archs for ROCm

tsrw2048 committed 13 hours ago

4e50ca6a

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

magaonka-amd committed 12 days ago

cbcc4e01

Update rules_ml_toolchain reference

alekstheod committed 26 days ago

b2e32655

Override HIP_VISIBLE_DEVICES per ROCm xdist worker

magaonka-amd committed 26 days ago

45b7af75

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

magaonka-amd committed 26 days ago

95aba756

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

psanal35 committed 26 days ago

acd00f93

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

magaonka-amd committed 26 days ago

a23d3768

Add a libtpu version guard in test_dce_sink_prevents_xla_dce

yashk2810 committed 41 days ago

a33ed614

Remove unused EXTRA_INDEX environment variable

nitins17 committed 41 days ago

79e8ef93

Allow shell globbing for libtpu wheel path in pytest_tpu.yml

nitins17 committed 41 days ago

83e0ac50

Prepare for JAX release 0.10.0

yashk2810 committed 41 days ago

08a228a9

Fix a bug in WeakValueInterner.

hawkinsp committed 41 days ago

2362cc76

Addressed an old pytype TODO

superbobry committed 41 days ago

c12e417e

Skip a test in OSS that relies on stdout

Rifur13 committed 41 days ago

1d0523d6

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

jburnim committed 41 days ago

01c0f46f

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

Google-ML-Automation committed 41 days ago

ad793d6a

Disable the eigh tests under asan.

hawkinsp committed 41 days ago

40e7c70d

USE `PYTHON_RUNFILES` to search for `libnvshmem_host.so.3`.

ybaturina committed 41 days ago

ca6e6f61

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

hawkinsp committed 41 days ago

e42b6a7c

[pallas:sc] Removed a few test skips

superbobry committed 41 days ago

e149dc87

Use uv-installed Python for pytest presubmits in place of setup-python one.

belitskiy committed 41 days ago

d80e79c2

Merge pull request #36810 from superbobry:pyrefly-2

Google-ML-Automation committed 41 days ago

73978577

[Mosaic GPU] Align `IsTransferableTmemRegisters` constraint with `TMEMRef.store` logic.

allanrenucci committed 41 days ago

c563d4a6

[Mosaic GPU] Make `alias_id` an optional property of `slice_smem`.

bchetioui committed 41 days ago

369d6368

[Pallas] Indicate in the quickstart guide that Mosaic GPU is the default (and maintained) backend on GPU.

bchetioui committed 41 days ago

1ea6146c

[Pallas:SC] Support get and swap with bitcast refs

apaszke committed 41 days ago

d37a62d7

[Mosaic:GPU] Fix test_multimem_store_contiguous_tma hang in OSS.

PatriosTheGreat committed 41 days ago

3af2f898

Older