ROCm/jax

Pull Requests Commits

[export] Handle AMD device kinds in get_platform_from_device_kind

zahiqbal committed 23 days ago

822f37f4

Merge pull request #37471 from TheSkyentist:main

Google-ML-Automation committed 23 days ago

033a21e2

Merge pull request #37589 from samanklesaria:fix_wald

Google-ML-Automation committed 23 days ago

501caca4

[Pallas] In TPU Interpret Mode, replace ordered IO callbacks with unordered

jburnim committed 23 days ago

fc936322

Add owens_t to jax.scipy.special

TheSkyentist committed 23 days ago

bc285cbd

Fix a problem where shared subjaxprs weren't printed as references if they were inside an outer jit.

hawkinsp committed 23 days ago

99ac7233

Merge pull request #37605 from gnecula:export_platform

Google-ML-Automation committed 23 days ago

b582ae4e

[Pallas:SC] Add a missing test skip

apaszke committed 23 days ago

635a7d03

Removed unnecessary Pyrefly suppressions

superbobry committed 23 days ago

db07004a

Patch rules_python to support local wheels

meteorcloudy committed 23 days ago

98098cc4

Append single_version_platform_override for custom Python toolchain

meteorcloudy committed 23 days ago

7ee66ab7

[Pallas:MGPU] Extend support for `get`/`swap` to untiled swizzled refs.

allanrenucci committed 23 days ago

0c27e7a4

Merge pull request #37616 from superbobry:maint

Google-ML-Automation committed 23 days ago

571f2a69

Bumped Pyrefly to 1.0.0

superbobry committed 23 days ago

f271a0da

[Pallas] Only pass `transpose_lhs_hint=True` for implicit transpose.

WindQAQ committed 23 days ago

7b92d791

[Pallas:SC] Add tests for constants of various types and shapes

apaszke committed 23 days ago

ca655d79

[Pallas/Mosaic GPU] Support lowering `memref.cast` on TMEM when using warpgroup semantics.

bchetioui committed 23 days ago

1f3025b3

[Pallas:MGPU] Raise when loading/storing from/to a non-transposed ref using a transposed layout.

allanrenucci committed 23 days ago

231b7351

[Mosaic GPU] Support lowering SMEM slices in a warp context.

bchetioui committed 23 days ago

876b1171

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

Google-ML-Automation committed 23 days ago

5de5427e

[Mosaic GPU] Allow `is_known_divisible` to go up `warp_map` arguments.

bchetioui committed 23 days ago

38b9dc9d

[export] Add serialization for AbstractDevice.platform

gnecula committed 23 days ago

d91cd88f

Merge pull request #37353 from Saumay:add-scipy-linalg-helmert

Google-ML-Automation committed 23 days ago

a574df9c

Fix export test on p100

yashk2810 committed 23 days ago

50e3d759

implement more rules for eval_jaxpr_p

Cristian Garcia committed 23 days ago

cf731a17

roll forward with fixes

mattjj committed 23 days ago

0ed2ac83

[pallas-mpmd] Remove `mpmd_meshes` from the lowering context in favor of meshes on avals.

rdyro committed 23 days ago

e6096e31

Track platform on AbstractDevice. We already track device_kind and platform is a coarser version i.e. we won't get new trace time cache misses (you can have multiple device kinds for 1 platform).

yashk2810 committed 24 days ago

e9c29d2f

Add jax.scipy.linalg.helmert

Saumay committed 24 days ago

8e1a04c6

Merge pull request #37580 from TheSkyentist:erfcx

Google-ML-Automation committed 24 days ago

05ed8208

Older