ROCm/jax

Pull Requests Commits

[ROCm][TEST-ONLY] mirror unversioned libhsa-runtime64.so to test RCCL workaround

gulsumgudukbay committed 10 hours ago

c14eb6c5

[ROCm] set HSA_NO_SCRATCH_RECLAIM=1 in test workflows

gulsumgudukbay committed 10 hours ago

f7773ab3

[ROCm][TEST-ONLY] strip /opt/rocm before tests to verify rpath resolution

gulsumgudukbay committed 10 hours ago

81b5c1f9

[ROCm] Bake TheRock-aware RUNPATH into rocm plugin/pjrt wheels

gulsumgudukbay committed 10 hours ago

22928460

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

mminutoli committed 11 hours ago

e3a41959

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 11 hours ago

a57dc0f9

Add org-based guards to all CI caller workflows

mminutoli committed 11 hours ago

7e9f2f20

Add __slots__ to dataclasses in jax.

hawkinsp committed 12 hours ago

e6848a3e

[NFC] Simplify `TiledLayout.remove_dimension`.

allanrenucci committed 12 hours ago

24cb778a

Add `ShapeDtypeStruct.like` classmethod. It does not work on Ref/AbstractRef inputs yet.

yashk2810 committed 13 hours ago

71fee995

[export] Expand the new test for typed prng keys

gnecula committed 13 hours ago

1b63c730

[Mosaic SC] Generalize broadcastSubelements to support SC shapes

apaszke committed 13 hours ago

ec89afbc

[pallas:sc] Adding float16 arrays no longer crashes

superbobry committed 15 hours ago

71f179a9

Reverts 8f711b489697b4172f17f0cfe47cf0f373f59bf3

chsigg committed 15 hours ago

69665e58

[pallas:mgpu] Removed an unnecessary test for `pl.pallas_call` monkey patching

superbobry committed 15 hours ago

cd5f6566

Update bug reference in test skip.

bchetioui committed 17 hours ago

ac0c1325

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

Google-ML-Automation committed 18 hours ago

0f8e113a

[Mosaic:GPU] Remove supports_cross_device_collectives since collective metadata currently support multi-process usage as well.

PatriosTheGreat committed 19 hours ago

f255b83f

Intern JaxprEqnContext

yashk2810 committed 1 day ago

5d6e541c

Use a weak_value_interner for LoweringCacheKey.

hawkinsp committed 1 day ago

864e6d52

Cache construction of JaxprEqnContext.

hawkinsp committed 1 day ago

89a67cb7

[Pallas TPU] Bugfix: ensure we physicalize output avals during lowering

sharadmv committed 1 day ago

eaebe4eb

Replace references to xla_client._xla with jax._src.lib._jax.

hawkinsp committed 1 day ago

e44e041c

Merge pull request #37564 from Saumay:add-scipy-special-boxcox

Google-ML-Automation committed 1 day ago

5799b549

Don't squeeze singleton dimension if their spec is not None in jnp.matmul.

yashk2810 committed 1 day ago

085710a8

Add a jax.extend.xla.register_hlo_module_transformation API.

hawkinsp committed 1 day ago

3a5e9c27

Roll back PR #37872; it led to timeouts in downstream projects.

Jake VanderPlas committed 1 day ago

92ae3541

Consume memory space + core type in isGather instead of MemRefs.

nvgrw committed 1 day ago

f3a724d8

[JAX][CPU] Skip running `testMathSpecialFloatValues` on ARM for `float16` `tanh`.

bchetioui committed 1 day ago

80df7b53

[pallas:mgpu] Deprecated using `pl.pallas_call` for Mosaic GPU kernels

superbobry committed 1 day ago

395f039c

Older