ROCm/jax

Pull Requests Commits

Add workflow for build wheels and running bazel tests on ROCm

charleshofer committed 1 day ago

bec3b790

Reimplement weakref_lru_cache on top of a custom hash map type.

hawkinsp committed 2 days ago

b2ccb32c

[Mosaic:GPU] Add test-case for several mosaic ops.

PatriosTheGreat committed 2 days ago

a25a24df

[Mosaic:GPU] Move collective kernel loading to the prepare stage.

PatriosTheGreat committed 2 days ago

f3bf01ad

[Mosaic:GPU] Re-use the barrier buffers between executions to prevent situation when the buffer were released before the end of cross-device barrier.

PatriosTheGreat committed 2 days ago

1609c18f

[mosaic_gpu] Fixed mosaic-gpu-resolve-trivial-locations

superbobry committed 2 days ago

e3adc6ad

[Mosaic TPU] Make sure to take the export flag into account when not using Pallas

apaszke committed 2 days ago

b9926c39

[Mosaic GPU] fix annotation typo: s/pytype/pylint/

cota committed 2 days ago

10a59a5e

[mosaic] Added pattern canonicalizing tpu.memref_squeeze(memref.cast(...))

superbobry committed 2 days ago

3ad94844

[pallas:sc] Moved `jaxpr_call` into core Pallas

superbobry committed 2 days ago

6400fa1a

Remove references to tsl/platform/logging.h header.

hawkinsp committed 2 days ago

89469861

Reverts ee4cea11e08b92f23f464be2e7795f5eb9fe9557

akuegel committed 2 days ago

143c2366

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

danielsuo committed 2 days ago

0d37f593

[Pallas][TPU kernel interpreter] Consistently use `pltu.HBM` to refer to HBM in the interpreter.

Google-ML-Automation committed 2 days ago

caf1324f

[Pallas:MGPU] Allow inline_mgpu inside a warp context

apaszke committed 2 days ago

df6415b4

Merge pull request #35102 from ROCm:set_release_rpaths_to_rocm_so_targets

Google-ML-Automation committed 2 days ago

717f99a4

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

Google-ML-Automation committed 2 days ago

08a51cea

Add a test for replicated -> sharded+unreduced reshard

yashk2810 committed 2 days ago

d506abdb

Add an experimental `top_level_all_gather` API (private for now).

yashk2810 committed 2 days ago

a6df1f29

Merge pull request #35422 from levskaya:pyrefly5

Google-ML-Automation committed 2 days ago

e008d9d8

[pmap] Accelerate deprecation of PmapSharding.

danielsuo committed 2 days ago

11af98b4

Merge pull request #35421 from levskaya:pyrefly4

Google-ML-Automation committed 2 days ago

019a87b1

[linux-arm64] Reduce parallelism to 32 for nogil pytest cpu to avoid OOMs.

danielsuo committed 2 days ago

d873c960

Merge pull request #35448 from jakevdp:discharge-hijax

Google-ML-Automation committed 2 days ago

153937a3

Merge pull request #35452 from jakevdp:fix-ad-checkpoint

Google-ML-Automation committed 2 days ago

c08a1707

Merge pull request #35219 from jakevdp:fix-tracer-tag

Google-ML-Automation committed 2 days ago

6d14b201

Fix undefined name in ad_checkpoint

jakevdp committed 2 days ago

e81437f2

Generalize how op verifiers checking issuing core

naummo committed 2 days ago

008467a7

[typing] fixed remaining scipy pyrefly typing issues

levskaya committed 2 days ago

47291b9f

[typing] miscellaneous pyrefly typing fixes

levskaya committed 2 days ago

85f6ef56

Older