ROCm/jax

Pull Requests Commits

Get device settings from file

charleshofer committed 17 days ago

2db1550d

Add GPU isolation check

charleshofer committed 20 days ago

16675a79

Error out if aval is HiType in `SDS.like`. `SDS.like` is not polymorphic. We need multiple things to make a polymorphic function work nicely.

yashk2810 committed 20 days ago

15a6dfd1

Add __slots__ to dataclasses in jax.

hawkinsp committed 20 days ago

e6848a3e

[NFC] Simplify `TiledLayout.remove_dimension`.

allanrenucci committed 20 days ago

24cb778a

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

yashk2810 committed 20 days ago

71fee995

[export] Expand the new test for typed prng keys

gnecula committed 20 days ago

1b63c730

[Mosaic SC] Generalize broadcastSubelements to support SC shapes

apaszke committed 20 days ago

ec89afbc

[pallas:sc] Adding float16 arrays no longer crashes

superbobry committed 20 days ago

71f179a9

Reverts 8f711b489697b4172f17f0cfe47cf0f373f59bf3

chsigg committed 20 days ago

69665e58

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

superbobry committed 20 days ago

cd5f6566

Update bug reference in test skip.

bchetioui committed 20 days ago

ac0c1325

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

Google-ML-Automation committed 20 days ago

0f8e113a

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

PatriosTheGreat committed 21 days ago

f255b83f

Intern JaxprEqnContext

yashk2810 committed 21 days ago

5d6e541c

Use a weak_value_interner for LoweringCacheKey.

hawkinsp committed 21 days ago

864e6d52

Cache construction of JaxprEqnContext.

hawkinsp committed 21 days ago

89a67cb7

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

sharadmv committed 21 days ago

eaebe4eb

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

hawkinsp committed 21 days ago

e44e041c

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

Google-ML-Automation committed 21 days ago

5799b549

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

yashk2810 committed 21 days ago

085710a8

Add a jax.extend.xla.register_hlo_module_transformation API.

hawkinsp committed 21 days ago

3a5e9c27

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

Jake VanderPlas committed 21 days ago

92ae3541

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

nvgrw committed 21 days ago

f3a724d8

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

bchetioui committed 21 days ago

80df7b53

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

superbobry committed 21 days ago

395f039c

[pallas:sc] Removed a few unnecessary test skips

superbobry committed 21 days ago

709095b9

[Mosaic GPU] Support slicing for non defaults TMEM layouts with 128 rows.

allanrenucci committed 21 days ago

80e49d9a

[pallas] Removed tests for cross-platform exports

superbobry committed 21 days ago

e59492a5

[pallas:sc] Do not skip SC tests on Cloud TPUs

superbobry committed 21 days ago

ae32a398

Older