ROCm/jax

Pull Requests Commits

Fix skip message: use correct function name testEighTinyNorm

AratiGanesh committed 20 days ago

383f9339

Skipping testEighTinyNorm due to hipSolver issues

AratiGanesh committed 24 days ago

560aea51

Increase shard count for gpu_pallas_interpret_test.

belitskiy committed 24 days ago

0f1e9c54

[logging] Avoid attaching the logging handler to jax for "NOTSET".

rdyro committed 24 days ago

63cd5913

Fix a Lowering issue for Pallas Kernel with Dynamic Shapes

Google-ML-Automation committed 24 days ago

4f82c7c3

[pallas:sc] `pltpu.sync_copy` now accepts `add=`

superbobry committed 25 days ago

21efe9a6

[pallas:sc] Added a primitive for atomic synchronous fetch-and-add in SMEM

superbobry committed 25 days ago

89fc1a93

Use std::map instead of std::unordered_map for CPU metadata buffers.

Google-ML-Automation committed 25 days ago

4fb45c8a

Add collective metadata version of collective matmul.

Google-ML-Automation committed 25 days ago

753298d3

[Pallas SC] Add support for `jax.empty_ref`.

bchetioui committed 25 days ago

1073b4c1

Fix crush due to non-trivial object destruction

Google-ML-Automation committed 25 days ago

00865dd5

[pallas:sc] Consistently use `math.prod` in tests

superbobry committed 25 days ago

9986cd87

Merge pull request #34871 from ezhulenev:device-groups

Google-ML-Automation committed 25 days ago

7bc8045f

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

Google-ML-Automation committed 25 days ago

9818cea0

[Mosaic:TPU][NFC] Move propagateTiledLayoutToConsumer to tpu_dialect.cc, expose in header

tlongeri committed 25 days ago

52ae55af

Allow reduce_sum to accept unreduced inputs. reduce_sum on a sharded input consists of 2 steps; (1) local reduction and (2) all-reduce across devices.

yashk2810 committed 25 days ago

0f262f22

Fix broadcast_in_dim_transpose rule

yashk2810 committed 25 days ago

b1631ce5

Add unreduced rule for split_p. Unreduced should just pass through since split is fully linear

yashk2810 committed 25 days ago

0b1e7179

Merge pull request #34827 from hrideymarwah15:fix/assert-false-to-valueerror

Google-ML-Automation committed 25 days ago

774b4b4d

[Mosaic TPU] Add support for bf16 sigmoid/logistic.

yueshengys committed 25 days ago

ca9500fa

[Mosaic-TPU] In tpu.memref_reshape verifier, also check target contiguity and use TiledLayoutAttr::tilesAreKnownContiguous

tlongeri committed 25 days ago

7f1ab37a

[Pallas] Remove bitwidth=32 condition when lowering cast.

WindQAQ committed 25 days ago

ba1c24de

Fix broadcast_in_dim tranpose rules if input is reduced, the ct should be unreduced.

yashk2810 committed 25 days ago

aa9e476a

Update JAX version regex to allow four-part versions.

belitskiy committed 25 days ago

ff35a774

[mosaic:gpu] Pass correct device groups when requesting a clique

ezhulenev committed 25 days ago

5e7e18d7

Merge pull request #34776 from jakevdp:index-generalization

Google-ML-Automation committed 25 days ago

7c9c08a4

[Mosaic:TPU] Make tpu.load and tpu.store's sublane_stride a DefaultValuedAttr instead of an OptionalAttr

tlongeri committed 25 days ago

36f475ec

Raise an error if indivisible shardings is returned by XLA.

yashk2810 committed 25 days ago

822e9a7c

Merge pull request #34641 from ROCm:add_support_to_dynamic_set_of_test_deps

Google-ML-Automation committed 25 days ago

9b90ff2a

Merge pull request #34450 from ROCm:upstream_rocm_ci

Google-ML-Automation committed 25 days ago

1627be1e

Older