ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

22426519

[Pallas:MGPU] Fix the implementation of WGMMA with transposed RHS

apaszke committed 2 years ago

bbcc3eef

[sharding_in_types] If out_aval.sharding is not None and the user specified out_sharding is None, concretize it with the device assignment available and add it to the final out_shardings that's used for lowering and compilation.

yashk2810 committed 2 years ago

2153de4c

[pallas] run_scoped now supports partial discharge.

cperivol committed 2 years ago

f8a3c036

Add a dialect for Mosaic GPU.

bchetioui committed 2 years ago

ade480ff

Disable breaking test-case

Google-ML-Automation committed 2 years ago

eba57480

[Pallas:MGPU] Fix docstring for commit_shared

apaszke committed 2 years ago

e138e8e4

nit: Use `frozen` dataclasses rather than `unsafe_hash`.

tomhennigan committed 2 years ago

86155561

[Pallas:MGPU] Force alignment of SMEM allocations to 1024 bytes

apaszke committed 2 years ago

40945648

[Pallas:MGPU] Add support for indexed refs to WGMMA

apaszke committed 2 years ago

0ee9531e

[Pallas:MGPU] Properly commute indexing with other transforms

apaszke committed 2 years ago

f2edc83a

Add `_sharding` argument to broadcasted_iota as a private parameter which only works under sharding_in_types mode.

yashk2810 committed 2 years ago

4db212d2

Allow simple host call that uses host tensor as parameter/result in

Google-ML-Automation committed 2 years ago

dd542630

Activate the FFI implementation of SVD on GPU.

dfm committed 2 years ago

8361eb58

[sharding_in_types] Add lax.transpose sharding propagation rule

yashk2810 committed 2 years ago

3e634d95

[sharding_in_types] Support jnp.array with sharding_in_types. When the input array has a sharding, propagate it through without dropping the sharding.

yashk2810 committed 2 years ago

57a95a77

[sharding_in_types] Add reduce max, integer_pow and standard_unop sharding rules

yashk2810 committed 2 years ago

5df4878a

[sharding_in_types] Add broadcast_in_dim rule.

yashk2810 committed 2 years ago

e92e1191

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

93389ab5

Merge pull request #24345 from phu0ngng:cuda_custom_call

Google-ML-Automation committed 2 years ago

919f7c86

[Pallas:MGPU] Add support for transforms in user-specified async copies

apaszke committed 2 years ago

2d78b172

Rewrite mosaic concat to support operand shapes that do not align with native shapes, Expand tests to cover multi operand, batch dim concat, etc.

Google-ML-Automation committed 2 years ago

6c2649fd

Add config option to log or fatal when jax.Arrays are GCed.

ICGog committed 2 years ago

ec279f9c

Fix breaking test-case

Google-ML-Automation committed 2 years ago

1b5cf5a4

[pallas:mosaic_gpu] Added `layout_cast`

superbobry committed 2 years ago

de7beb91

[Pallas:MGPU] Add lowerings for more ops

apaszke committed 2 years ago

0519db15

[Pallas:MGPU] Expose mgpu.commit_shared

apaszke committed 2 years ago

83b61cc9

[Pallas:MGPU] Add support for debug_print of arrays that use the WGMMA layout

apaszke committed 2 years ago

f72376ae

[Mosaic GPU] Add support for launching multiple warpgroups using core_map

apaszke committed 2 years ago

ef361f05

Merge pull request #24300 from ROCm:ci_rocm_readme

Google-ML-Automation committed 2 years ago

3bdc57dd

Newer Older