ROCm/jax

Pull Requests Commits

Avoid using array operations within lax.py operations.

dfm committed 332 days ago

4eada560

[pallas:mosaic_gpu] Added support for math functions to the WG lowering

superbobry committed 333 days ago

91340ea0

Update XLA dependency to use revision

Google-ML-Automation committed 333 days ago

f906d2b6

[Pallas/Mosaic GPU] Add lowerings/layout inference for all the necessary conversion ops when using Warpgroup semantics.

bchetioui committed 333 days ago

75d87020

Update stateful-computations.md

garymm committed 333 days ago

Verified 6a718b76

Update XLA dependency to use revision

Google-ML-Automation committed 334 days ago

b9fb69d1

Merge pull request #27006 from dfm:more-direct-lin-debug-info

Google-ML-Automation committed 334 days ago

922935a9

A few more fixes for debug_info tests with direct_linearize.

dfm committed 335 days ago

36d515ed

Update XLA dependency to use revision

Google-ML-Automation committed 335 days ago

04696b4d

[Mosaic TPU][Pallas] Add pl.reciprocal

bythew3i committed 335 days ago

0f0636af

Merge pull request #27010 from mattjj:direct-linearize-fixes-3

Google-ML-Automation committed 335 days ago

4988adcc

[direct-linearize] fix name_stack bugs

mattjj committed 335 days ago

fe26c19b

Merge pull request #27005 from mattjj:direct-linearize-fixes-2

Google-ML-Automation committed 335 days ago

4660d7b6

fixups that we meant to include in #26427

mattjj committed 335 days ago

251b93eb

Support MHA in ragged paged attention for packed type

bythew3i committed 335 days ago

041f5757

Merge pull request #26427 from mattjj:direct-linearize-fixes

Google-ML-Automation committed 335 days ago

6095af05

Merge pull request #27001 from mattjj:yash-scan

Google-ML-Automation committed 335 days ago

d8497796

Merge pull request #26979 from mattjj:26936

Google-ML-Automation committed 335 days ago

1870176e

[scan] when num_trips==0, don't generate weird size-zero reshapes

mattjj committed 335 days ago

f4f31f89

shard_map and other fixes to direct-linearize

mattjj committed 335 days ago

7c2f8423

[mutable-arrays] read values should have the same explicit sharding as ref

mattjj committed 335 days ago

0e30a3ac

[JAX] Fix api_benchmark broken by https://github.com/jax-ml/jax/pull/26569

hyeontaek committed 335 days ago

17827886

Add the len(arg) to the error message for static_argnums

Google-ML-Automation committed 335 days ago

ccf72782

[sharding_in_types] Fix a bug where `empty_array` in scan was created with the wrong spec when `unroll > 1`.

yashk2810 committed 335 days ago

9f37b519

[mgpu] Debug print arrays.

cperivol committed 335 days ago

eeccc67c

[Mosaic GPU] Add a better explanation for the transposed layout

apaszke committed 335 days ago

1bef8b61

[Mosaic TPU] Enable all conversions involving fp8 types on TPUv5+

apaszke committed 335 days ago

40238929

[pallas:mosaic_gpu] `copy_smem_to_gmem` now allows skipping `cp.async.commit_group`

superbobry committed 335 days ago

928caf83

[Mosaic GPU] Add a new layout to help with transposing WGMMA results

apaszke committed 335 days ago

65462fe6

Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings.

yashk2810 committed 335 days ago

f8b98993

Newer Older