ROCm/jax

Pull Requests Commits

Removed unused `jax_remat_opt_barrier` config option

superbobry committed 327 days ago

ff00fa91

Apply forwarding in pjit linearization rule to avoid intermediate copies.

dfm committed 327 days ago

dbc3bcd3

Deprecate public export of mlir.custom_call.

dfm committed 327 days ago

5a3fc606

Reverts 735cec18cb2f8dff2aea5e503fd886a37aee094e

danielsuo committed 327 days ago

a099b285

Temporarily skip JaxNumpyErrorTests in multi-thread environments

ayaka14732 committed 327 days ago

412f88e2

Remove accidental exports jax.interpreters.mlir.{hlo,func_dialect}.

hawkinsp committed 327 days ago

70485e31

[Mosaic GPU] Add missing to/from tiled layout attributes with replicated lane dimensions.

dimitar-asenov committed 327 days ago

83572e17

Update XLA dependency to use revision

Google-ML-Automation committed 327 days ago

5c0f8858

Merge pull request #27773 from hawkinsp:race

Google-ML-Automation committed 327 days ago

724b13ac

`jex.core.Var` is no longer ordered

superbobry committed 328 days ago

c2aa811c

Add more TSAN skips to avoid timeouts

apaszke committed 328 days ago

153fa228

Add a missing jaxlib version check in Pallas TPU lowering

apaszke committed 328 days ago

4596ee3c

[Mosaic GPU] Delete mentions of `WGMMARowFragLayout` in `layouts.py`.

bchetioui committed 328 days ago

9b850a9e

Removed unused deprecations

superbobry committed 328 days ago

6e93fa34

Fix some test timeouts

apaszke committed 328 days ago

075d88fe

Fix a race in pjit under free threading.

hawkinsp committed 328 days ago

695ee8f3

Use `contextlib.nullcontext` instead of `trivial_ctx`

superbobry committed 328 days ago

245194ff

[Mosaic GPU] Support Slice and Transpose in the Pallas WGMMA lowering

dimitar-asenov committed 328 days ago

90cfa99a

Raise an error if the type passed to `axis_types` argument of `Mesh` and `AbstractMesh` is not `jax.sharding.AxisType`.

yashk2810 committed 328 days ago

cccc34dc

Fix deadlock when computing cached Sharding::type() values.

hawkinsp committed 328 days ago

8a6efa31

Automated Code Change

Google-ML-Automation committed 328 days ago

7874d79f

Automated Code Change

Google-ML-Automation committed 328 days ago

d1c7ba43

Automated Code Change

Google-ML-Automation committed 328 days ago

3f083cae

Automated Code Change

Google-ML-Automation committed 328 days ago

477b1082

Automated Code Change

Google-ML-Automation committed 328 days ago

ad36f7f2

Update XLA dependency to use revision

Google-ML-Automation committed 329 days ago

b1b54f9b

Update XLA dependency to use revision

Google-ML-Automation committed 329 days ago

2e62693f

Check that memory_kind of an aval is always None

yashk2810 committed 330 days ago

fc5d9a4f

[pallas:mosaic_gpu] Fixed a typo in `_barrier_arrive_pp_eqn`

superbobry committed 330 days ago

aab66139

[Mosaic GPU] Limit the maximum number of registers per thread to 255.

Rifur13 committed 330 days ago

d81c0ffe

Newer Older