ROCm/jax

Pull Requests Commits

[XLA:Python] [JAX] Move HloPass bindings to XLA/Python.

hawkinsp committed 1 year ago

e1ca6aba

Block yt_kernel in test_upcast_to_wgmma before the assert

Google-ML-Automation committed 1 year ago

d58a7346

Add block_until_ready to buffer_callback tests where needed.

dfm committed 1 year ago

f5d621f4

Merge pull request #28500 from hawkinsp:tsan

Google-ML-Automation committed 1 year ago

ecf41a74

Merge pull request #28498 from kirillbobyrev:fix-dlpack-docs-formatting

Google-ML-Automation committed 1 year ago

e66c55aa

Improve TSAN CI build.

hawkinsp committed 1 year ago

6fa05143

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

55fb40f5

Fix `dlpack` docs formatting

kirillbobyrev committed 1 year ago

Verified b5db71bd

Fix with_sharding_constraint when a prng key is passed to it. The problem was that the input aval's rank and the mlir input node's rank did not match. Converting the input aval to physical aval fixes that.

yashk2810 committed 1 year ago

a96e175c

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c15f26a1

[Pallas][Mosaic GPU] Add support for blackwell mma instruction.

justinjfu committed 1 year ago

43b784c3

[IFRT] Introduce `xla::ifrt::ShardingRef` type alias

hyeontaek committed 1 year ago

9fef4001

Put the Box, List handling inside _flatten_boxes

yashk2810 committed 1 year ago

b93ffa70

Merge pull request #28489 from rickeylev:disable.legacy.runfiles

Google-ML-Automation committed 1 year ago

b7595d0b

Add Bazel CPU tests with `py_import` dependency to continuous tests.

Google-ML-Automation committed 1 year ago

22375ab5

[pallas:mgpu] Allow WGMMA accuculator to cond yielded values.

cperivol committed 1 year ago

5a3605d3

chore: disable legacy external runfiles

rickeylev committed 1 year ago

653085d4

fix breakage from #28318

mattjj committed 1 year ago

c7f3d1c0

Merge pull request #28481 from hawkinsp:l4

Google-ML-Automation committed 1 year ago

ffd1f616

Merge pull request #28429 from hawkinsp:requirements

Google-ML-Automation committed 1 year ago

5116dd52

Explicitly fail on pallas w/ disable_jit until we can decide if/how we want to support it.

Google-ML-Automation committed 1 year ago

5efbaa8d

fix breakage from #28318 due to colliding kwargs

mattjj committed 1 year ago

832f3bac

Merge pull request #28283 from yhtang:pr-add-k8s-retry

Google-ML-Automation committed 1 year ago

9c5e03ff

[pallas:mgpu] Add select op support to divisibility inferencece.

cperivol committed 1 year ago

06161ef8

Merge pull request #28482 from jakevdp:fix-docs

Google-ML-Automation committed 1 year ago

0454a754

Mark `out_sharding` as a keyword argument in the `pyi` files too for `jnp.ravel`, `jnp.reshape` and `jnp.dot`

yashk2810 committed 1 year ago

bf5d1c00

DOC: hide config options TOC on front page

jakevdp committed 1 year ago

2a90dae6

Update ManualComputation round-tripping to not have round-tripping attrs on the CallOps.

bartchr808 committed 1 year ago

f4626f42

Skip Pallas tests that are failing in CI on L4 GPUs.

hawkinsp committed 1 year ago

b26d8553

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c8f7e39f

Older