ROCm/jax

Pull Requests Commits

alekstheod committed 33 days ago

7d55b300

Change docker img

alekstheod committed 33 days ago

39b02bc1

Include released wheels testing pipeline

alekstheod committed 33 days ago

98aeb33b

Use python version

alekstheod committed 33 days ago

cbb3f054

Include prebuilt rocm plugin pkgs as deps to tests

alekstheod committed 33 days ago

51e574a5

Add workflow for build wheels and running bazel tests on ROCm

charleshofer committed 36 days ago

bec3b790

Reimplement weakref_lru_cache on top of a custom hash map type.

hawkinsp committed 36 days ago

b2ccb32c

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

PatriosTheGreat committed 36 days ago

a25a24df

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

PatriosTheGreat committed 36 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 36 days ago

1609c18f

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

superbobry committed 36 days ago

e3adc6ad

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

apaszke committed 36 days ago

b9926c39

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

cota committed 36 days ago

10a59a5e

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

superbobry committed 36 days ago

3ad94844

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

superbobry committed 36 days ago

6400fa1a

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

hawkinsp committed 36 days ago

89469861

Reverts ee4cea11e08b92f23f464be2e7795f5eb9fe9557

akuegel committed 36 days ago

143c2366

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

danielsuo committed 36 days ago

0d37f593

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

Google-ML-Automation committed 36 days ago

caf1324f

[Pallas:MGPU] Allow inline_mgpu inside a warp context

apaszke committed 36 days ago

df6415b4

Merge pull request #35102 from ROCm:set_release_rpaths_to_rocm_so_targets

Google-ML-Automation committed 36 days ago

717f99a4

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

Google-ML-Automation committed 36 days ago

08a51cea

Add a test for replicated -> sharded+unreduced reshard

yashk2810 committed 37 days ago

d506abdb

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

yashk2810 committed 37 days ago

a6df1f29

Merge pull request #35422 from levskaya:pyrefly5

Google-ML-Automation committed 37 days ago

e008d9d8

[pmap] Accelerate deprecation of PmapSharding.

danielsuo committed 37 days ago

11af98b4

Merge pull request #35421 from levskaya:pyrefly4

Google-ML-Automation committed 37 days ago

019a87b1

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

danielsuo committed 37 days ago

d873c960

Merge pull request #35448 from jakevdp:discharge-hijax

Google-ML-Automation committed 37 days ago

153937a3

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

Google-ML-Automation committed 37 days ago

c08a1707

Older