Commits ROCm/jax - SemanticDiff

Fix plugin reference for ROCm >= 7.0.

mminutoli committed 186 days ago

579d387c

[Mosaic GPU][NFC] Move tests that don't require a GPU to platform agnostic tests.

allanrenucci committed 186 days ago

1df8d4cf

[Pallas] Better error message for grid/dimension semantics mismatch.

justinjfu committed 186 days ago

e7545eca

[Pallas:MGPU] Add initial support for TMEM under WG semantic.

allanrenucci committed 186 days ago

a8592fe0

[Mosaic TPU] No longer implicitly trigger single-buffering when running out of VMEM

apaszke committed 186 days ago

588c2812

[pmap] Disable VmapPmapCollectivesTest.testCollectivesWithVmap2 for 4x gpu tests.

danielsuo committed 186 days ago

42365834

[Mosaic GPU][NFC] Simplify `VerifyTmemRefType` by getting the context from the operation.

allanrenucci committed 186 days ago

150c336c

[Mosaic GPU][NFC] Consolidate `is_signed` derivation in dialect lowering.

allanrenucci committed 186 days ago

8bda0065

[Mosaic GPU][NFC] Restrict valid `packing` values.

allanrenucci committed 186 days ago

46d427b6

Remove ExecuteShardedArg.

hawkinsp committed 186 days ago

dc962146

[Mosaic GPU] Update `TODO` to indicate the root cause of numerical failures.

bchetioui committed 186 days ago

98ca8924

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

Google-ML-Automation committed 186 days ago

85bba835

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

Google-ML-Automation committed 187 days ago

6537c437

Merge pull request #33021 from cgarciae:hijax-coverage-tests

Google-ML-Automation committed 188 days ago

f386951c

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

Google-ML-Automation committed 188 days ago

b11d86c8

[jax:ffi] Declare ffi::TypeInfo as a struct static member

ezhulenev committed 189 days ago

4f88b028

[Mosaic TPU] Allow small sublane tilings in second/third minor transpose for unpacked dtype instead of forcing native tiling, which can be more efficient when the second minor dimension of input tensor is small.

yueshengys committed 189 days ago

af44ec5c

[xla:ffi] Migrate to xla::ffi::MakeTypeInfo() API

ezhulenev committed 189 days ago

c5c103cd

Reverts f72308e83aabcb061b16d58fc8330587de3cee02

hawkinsp committed 189 days ago

437d260f

add grad hijax coverage tests

cgarciae committed 189 days ago

709e3364

Merge pull request #33041 from mattjj:hijax-vjp-accums

Google-ML-Automation committed 189 days ago

f6edb899

Pass git hash to artifact build for release and nightly builds

nitins17 committed 189 days ago

40e9dc9d

[hijax] upgrade VJPHiPrimitive.def_bwd to work with accums

mattjj committed 189 days ago

eff56af4

[CI] Reduce the number of concurrent test jobs in the TSAN workflow.

belitskiy committed 189 days ago

08c63d38

[Mosaic GPU][NFC] Replace calls to `llvm.getelementptr` with `utils.getelementptr`.

allanrenucci committed 189 days ago

0a972ed4

[Mosaic GPU] Remove `TraversalOrder` from `traverse_op`.

bchetioui committed 189 days ago

7239d15f

[Mosaic TPU] Allow using replicated lane offset for intermediate safe offsets in mask relayout since relayout like {0,*},(8,128) -> {0,*},(4,128) has been just supported.

yueshengys committed 189 days ago

dc7d56e6

[Pallas/Mosaic GPU] Add support for `s8` WGMMA with lhs in registers.

bchetioui committed 189 days ago

7a007ea9

[Mosaic GPU][NFC] Silence some linter warning.

bchetioui committed 189 days ago

14d89b51

Repurpose `tpu_sparsecore` annotation for offloading purposes

vsytch committed 189 days ago

2d1479dc