ROCm/jax

Pull Requests Commits

[ROCm] Add release pipeline CI workflows

mminutoli committed 1 day ago

f25f2b47

CI: Remove --device=/dev/dri from pytest_rocm container options

magaonka-amd committed 1 day ago

60bde38f

[ROCm] Update CI runner labels to amd-do-linux.jax.gpu.gfx950.{1,4,8} (#807)

mminutoli committed 1 day ago

a1f4b433

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

gulsumgudukbay committed 1 day ago

fcb64d1c

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

gulsumgudukbay committed 1 day ago

f6eec847

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

mminutoli committed 1 day ago

00cc9f67

Make sure we add unreduced constraints at the input/output boundaries of `kCall`.

yashk2810 committed 1 day ago

bbd2b246

[Pallas] In GPU interpret mode, implement race detection for async TMA operations

Paul Biberstein committed 1 day ago

da78585e

Add DlPack for host memory on TPU.

pschuh committed 1 day ago

51a20c5f

Remove stale reference to tsan_suppressions for 3.13t which we have dropped.

hawkinsp committed 1 day ago

6a19c8b5

Skip mosaic test under pytest that fails with empty ptx capture.

hawkinsp committed 1 day ago

0c976f84

[pallas:mosaic] Use `contextlib.redirect_stdout` instead of a custom helper

superbobry committed 1 day ago

2a320738

Merge pull request #38663 from ROCm:jd/upstream-rocm-install-docs

Google-ML-Automation committed 1 day ago

526a94c2

[jaxlib] Fix memory leak when PJRT plugins lack optional custom partitioner extensions

jenriver committed 1 day ago

4f45f261

Default core_id_mutability to False when core_ids is specified.

Google-ML-Automation committed 1 day ago

2adbd8ef

[Mosaic TPU] Support dynamic rotate on the second minor dim with zero stride for packed types, we require the dim size to be divisible by `packing x sublane_count`.

yueshengys committed 1 day ago

7c494973

Convert all usages of gpuError_t to use absl::Status consistently (or convert).

pschuh committed 1 day ago

11b55240

Pad the partitions with None if it's length is less than the `axis` in get_sharding_for_vmap

yashk2810 committed 1 day ago

20a74963

[pallas:mosaic] `pl.debug_print` now supports 32-bit floats

superbobry committed 1 day ago

64a2fcfb

Reverts 776018ae0b4f2972f15cd23673fbca817a267f7d

zvikinoza committed 1 day ago

bff1d544

Use pspec.partitions in _make_lengths_same

Google-ML-Automation committed 1 day ago

e7c30496

Expose `jax.ffi.abstract_token`

Jake VanderPlas committed 1 day ago

a6c992da

[Pallas:MGPU] Add Lane semantics lowering for `lax.concatenate`.

allanrenucci committed 1 day ago

67cae6e2

[mosaic] Enabled `use-nameloc-as-prefix` for IR dumps

superbobry committed 1 day ago

055a433b

[Mosaic GPU] Add support for concatenating `FragmentedArray`s.

allanrenucci committed 1 day ago

425c9c03

Add PDL kernel launch integration to pallas mosaic gpu

rajaanand-g committed 1 day ago

5c65d574

Fix bug where jnp.int64(1 << 40) would truncate even if explicit x64 types are allowed.

hawkinsp committed 1 day ago

f7337b26

Merge pull request #38694 from aybchan:aybchan/skip-tile-scatter4

Google-ML-Automation committed 1 day ago

5f8829ca

Merge pull request #38724 from vfdev-5:pallas-fix-_gmem_slice_and_predicate

Google-ML-Automation committed 1 day ago

3673b557

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

Google-ML-Automation committed 1 day ago

4c290b69

Older