ROCm/jax

Pull Requests Commits

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

hawkinsp committed 315 days ago

e1ca6aba

Block yt_kernel in test_upcast_to_wgmma before the assert

Google-ML-Automation committed 315 days ago

d58a7346

Add block_until_ready to buffer_callback tests where needed.

dfm committed 316 days ago

f5d621f4

Merge pull request #28500 from hawkinsp:tsan

Google-ML-Automation committed 316 days ago

ecf41a74

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

Google-ML-Automation committed 316 days ago

e66c55aa

Improve TSAN CI build.

hawkinsp committed 316 days ago

6fa05143

Update XLA dependency to use revision

Google-ML-Automation committed 316 days ago

55fb40f5

Fix `dlpack` docs formatting

kirillbobyrev committed 316 days 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 316 days ago

a96e175c

Update XLA dependency to use revision

Google-ML-Automation committed 317 days ago

c15f26a1

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

justinjfu committed 317 days ago

43b784c3

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

hyeontaek committed 317 days ago

9fef4001

Put the Box, List handling inside _flatten_boxes

yashk2810 committed 317 days ago

b93ffa70

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

Google-ML-Automation committed 317 days ago

b7595d0b

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

Google-ML-Automation committed 317 days ago

22375ab5

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

cperivol committed 317 days ago

5a3605d3

chore: disable legacy external runfiles

rickeylev committed 317 days ago

653085d4

fix breakage from #28318

mattjj committed 317 days ago

c7f3d1c0

Merge pull request #28481 from hawkinsp:l4

Google-ML-Automation committed 317 days ago

ffd1f616

Merge pull request #28429 from hawkinsp:requirements

Google-ML-Automation committed 318 days ago

5116dd52

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

Google-ML-Automation committed 318 days ago

5efbaa8d

fix breakage from #28318 due to colliding kwargs

mattjj committed 318 days ago

832f3bac

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

Google-ML-Automation committed 318 days ago

9c5e03ff

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

cperivol committed 318 days ago

06161ef8

Merge pull request #28482 from jakevdp:fix-docs

Google-ML-Automation committed 318 days ago

0454a754

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

yashk2810 committed 318 days ago

bf5d1c00

DOC: hide config options TOC on front page

jakevdp committed 318 days ago

2a90dae6

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

bartchr808 committed 318 days ago

f4626f42

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

hawkinsp committed 318 days ago

b26d8553

Update XLA dependency to use revision

Google-ML-Automation committed 318 days ago

c8f7e39f

Older