Commits ROCm/jax - SemanticDiff

Remove old ROCm build code

charleshofer committed 105 days ago

21d69b19

Merge pull request #32982 from jakevdp:config-help

Google-ML-Automation committed 105 days ago

40b4d263

prepare tests for deprecation of `jax.scipy.special.sph_harm`

Jake VanderPlas committed 105 days ago

0bc1bacc

Support overrideable jax coordinator port and modify default port for Cloud TPU jax.distributed.initialize()

Google-ML-Automation committed 105 days ago

6ea9d94f

Use a faster pathway for i4 -> f8 canonicalization on TPU 7x: i4 -> i8 -> f8 instead of i4 -> i32 -> f32 -> f8.

18praveenb committed 105 days ago

9c9626c1

Remove the redundant context manager in lower_fun. We already enter into eqn's context one level above in jaxpr_subcomp

yashk2810 committed 105 days ago

44c13598

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

allanrenucci committed 105 days ago

1df8d4cf

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

justinjfu committed 105 days ago

e7545eca

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

allanrenucci committed 105 days ago

a8592fe0

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

apaszke committed 105 days ago

588c2812

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

danielsuo committed 105 days ago

42365834

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

allanrenucci committed 105 days ago

150c336c

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

allanrenucci committed 106 days ago

8bda0065

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

allanrenucci committed 106 days ago

46d427b6

Remove ExecuteShardedArg.

hawkinsp committed 106 days ago

dc962146

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

bchetioui committed 106 days ago

98ca8924

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

Google-ML-Automation committed 106 days ago

85bba835

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

Google-ML-Automation committed 107 days ago

6537c437

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

Google-ML-Automation committed 107 days ago

f386951c

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

Google-ML-Automation committed 108 days ago

b11d86c8

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

ezhulenev committed 108 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 108 days ago

af44ec5c

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

ezhulenev committed 108 days ago

c5c103cd

Reverts f72308e83aabcb061b16d58fc8330587de3cee02

hawkinsp committed 108 days ago

437d260f

add grad hijax coverage tests

cgarciae committed 108 days ago

709e3364

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

Google-ML-Automation committed 108 days ago

f6edb899

Pass git hash to artifact build for release and nightly builds

nitins17 committed 108 days ago

40e9dc9d

[hijax] upgrade VJPHiPrimitive.def_bwd to work with accums

mattjj committed 108 days ago

eff56af4

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

belitskiy committed 108 days ago

08c63d38

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

allanrenucci committed 108 days ago

0a972ed4