ROCm/jax

Pull Requests Commits

#pallas Avoid scaling the pallas_call CostEstimate by the batching dimension size if it is a dynamic shape.

Google-ML-Automation committed 297 days ago

9dcb93f0

[Pallas:MGPU] Remove a test skip for casts from 32-bit

apaszke committed 297 days ago

fd30312a

[Mosaic GPU] Remove result from `mgpu.InitializeBarrierOp`.

allanrenucci committed 297 days ago

9f5ca4fa

[Mosaic GPU] Add a `fuel` argument to `infer_layout`.

bchetioui committed 297 days ago

6cd20d59

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

Google-ML-Automation committed 297 days ago

675df7fe

Reverts 1781d2892759cf5565837b889ea48034830116e3

hawkinsp committed 298 days ago

66924519

Add a missing mesh_utils TPU kind.

hawkinsp committed 298 days ago

1781d289

Add a mirror for `rules_ml_toolchain` due to the error:

yuriivcs committed 298 days ago

63bd74e8

Dtype in transpose rule of gather should follow tangent rather that primal.

yliu120 committed 298 days ago

9f426e23

Delete old Shard Map import code path that isn't needed anymore.

bartchr808 committed 298 days ago

1c9263dd

[JaxMonitoring] Surface unregister listener functions as external APIs.

jshin1394 committed 298 days ago

ff815833

Consider this example:

yashk2810 committed 298 days ago

f59e61a2

Merge pull request #32982 from jakevdp:config-help

Google-ML-Automation committed 298 days ago

40b4d263

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

Jake VanderPlas committed 298 days ago

0bc1bacc

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

Google-ML-Automation committed 298 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 298 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 298 days ago

44c13598

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

allanrenucci committed 298 days ago

1df8d4cf

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

justinjfu committed 298 days ago

e7545eca

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

allanrenucci committed 298 days ago

a8592fe0

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

apaszke committed 298 days ago

588c2812

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

danielsuo committed 298 days ago

42365834

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

allanrenucci committed 298 days ago

150c336c

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

allanrenucci committed 298 days ago

8bda0065

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

allanrenucci committed 298 days ago

46d427b6

Remove ExecuteShardedArg.

hawkinsp committed 298 days ago

dc962146

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

bchetioui committed 298 days ago

98ca8924

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

Google-ML-Automation committed 298 days ago

85bba835

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

Google-ML-Automation committed 299 days ago

6537c437

[mutable-arrays] make one-level remat work with refs

mattjj committed 300 days ago

ddc67f06

Newer Older