ROCm/jax

Pull Requests Commits

[ROCm] Enable HLO module transform registration for GPU backends

mminutoli committed 21 days ago

5d1c0bc9

default is_high generalization

Cristian Garcia committed 21 days ago

c2591b29

Couple of changes:

yashk2810 committed 21 days ago

3b9d8b92

Release the GIL during portable MLIR serialization and deserialization.

hawkinsp committed 21 days ago

13bd7c3a

[Pallas/Mosaic] Add subcore id as an arg to `enqueue_dma`, `semaphore_signal` and `wait_dma` ops.

emilyfertig committed 21 days ago

2f479087

[Pallas/Mosaic GPU] Fix `cond` lowering under Lane semantics to use `FragmentedArray`s. consistently

bchetioui committed 21 days ago

7efe39e0

Merge pull request #38106 from jakevdp:heaviside-nan

Google-ML-Automation committed 21 days ago

66588223

Merge pull request #38124 from gnecula:hijax_vmap

Google-ML-Automation committed 21 days ago

8a0456ae

Update hypothesis test util to disable sharding after a failure (in shrinking phase) so that it can track down the failure correctly and not get confused by skipped tests.

18praveenb committed 21 days ago

4676543d

[mosaic_gpu] Support TMA_INDICES_4_LAYOUT for TMA gather/scatter (replicated).

cperivol committed 21 days ago

caa0b9ef

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 21 days ago

15a6dfd1

Add __slots__ to dataclasses in jax.

hawkinsp committed 21 days ago

e6848a3e

[NFC] Simplify `TiledLayout.remove_dimension`.

allanrenucci committed 21 days ago

24cb778a

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

yashk2810 committed 21 days ago

71fee995

[export] Expand the new test for typed prng keys

gnecula committed 21 days ago

1b63c730

[Mosaic SC] Generalize broadcastSubelements to support SC shapes

apaszke committed 21 days ago

ec89afbc

[hijax] Fix nested vmap for hi types

gnecula committed 21 days ago

c656efb5

[pallas:sc] Adding float16 arrays no longer crashes

superbobry committed 21 days ago

71f179a9

Reverts 8f711b489697b4172f17f0cfe47cf0f373f59bf3

chsigg committed 21 days ago

69665e58

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

superbobry committed 21 days ago

cd5f6566

Update bug reference in test skip.

bchetioui committed 22 days ago

ac0c1325

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

Google-ML-Automation committed 22 days ago

0f8e113a

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

PatriosTheGreat committed 22 days ago

f255b83f

Intern JaxprEqnContext

yashk2810 committed 22 days ago

5d6e541c

Use a weak_value_interner for LoweringCacheKey.

hawkinsp committed 22 days ago

864e6d52

Cache construction of JaxprEqnContext.

hawkinsp committed 22 days ago

89a67cb7

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

sharadmv committed 22 days ago

eaebe4eb

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

hawkinsp committed 22 days ago

e44e041c

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

Google-ML-Automation committed 22 days ago

5799b549

[bug] fix jnp.heaviside for NaN inputs

jakevdp committed 22 days ago

88b3d133

Older