ROCm/jax

Pull Requests Commits

Merge pull request #35712 from jakevdp:pyrefly-main-fix

Google-ML-Automation committed 109 days ago

4a592c9d

[ci] silence pyrefly error at head

jakevdp committed 109 days ago

d0ce0102

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

Google-ML-Automation committed 110 days ago

fc393eca

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

Google-ML-Automation committed 111 days ago

7b6360dd

[Mosaic] Add "add" attribute to tpu.store

tlongeri committed 111 days ago

23bede0c

[Mosaic:TPU] In tpu.memref_reshape, don't look through tpu.erase_layout and support empty tiles

tlongeri committed 111 days ago

d1495d28

Reverts a1ff74b306dc37e979bee7c8b56e1d9eb425bcce

gnecula committed 111 days ago

cfb68889

[Mosaic] Delete isContiguousMemRef in favor of TiledLayoutAttr::tilesAreKnownContiguous

tlongeri committed 111 days ago

2130ed7f

Use `ct_zero_from_primal` in shard_map_transpose to create zeros from residuals

yashk2810 committed 111 days ago

82f4db2a

Reverts c43e68919106c28b379ec9353439809e2ce7f797

gnecula committed 111 days ago

ab202431

[Mosaic] Clean up some logic to look past tpu.erase_layout

tlongeri committed 111 days ago

dedbc60d

Show detailed mismatch key paths in `flatten_axes()` tree prefix errors.

ashishenoyp committed 111 days ago

caee8a92

Error out if sorted dimension is sharded because that almost always entails an AG. It's better to reshard at the user level instead.

yashk2810 committed 111 days ago

acbca16b

Merge pull request #35111 from ROCm:rocm-eigh-export-compat

Google-ML-Automation committed 111 days ago

e578b6d4

Merge pull request #35521 from jax-ml:dependabot/pip/pydata-sphinx-theme-0.16.1

Google-ML-Automation committed 111 days ago

1808f592

[Mosaic] Mark tpu.store and tpu.load base as memref

tlongeri committed 111 days ago

27650dfb

Bump pydata-sphinx-theme from 0.14.4 to 0.16.1

dependabot[bot] committed 111 days ago

Verified 2fee16fa

Merge pull request #35689 from jakevdp:pyrefly-pallas

Google-ML-Automation committed 111 days ago

b69238e7

PR #35550: [ROCm] Add ROCm wheel build and test pipeline to continuous CI

alekstheod committed 111 days ago

9a0eddda

Set CUDA UMD driver version to 13.0.2 in all JAX bazel jobs.

ybaturina committed 111 days ago

80f3cf17

[pyrefly] type check jax/experimental/pallas

jakevdp committed 111 days ago

323ec0b8

Fix flaky test in ffi_test.

hawkinsp committed 111 days ago

6f3c7794

Skip multicast-based tests for when multicast is not yet implemented.

belitskiy committed 111 days ago

00b08ccb

Update HERMETIC_CUDA_UMD_VERSION 13.0.0 -> 13.0.2.

ybaturina committed 111 days ago

c80435cb

Merge pull request #35677 from gnecula:ser_abs_device

Google-ML-Automation committed 111 days ago

c43e6891

[export] Add support for serializating AbstractMesh.abstract_device

gnecula committed 111 days ago

5498417d

[Mosaic TPU] Add a rule to remove no-op tpu.memref_slice

apaszke committed 111 days ago

b7e106f4

[Mosaic GPU] Add `a_sparse_metadata` argument to `tcgen05_mma` dialect definition.

bchetioui committed 111 days ago

15a6469e

[Mosaic GPU] Add a `rank` field to the `Reduce` constraint.

bchetioui committed 111 days ago

8f232c94

[Mosaic GPU] Validate `source_dimensions` in `broadcast_in_dim`.

allanrenucci committed 111 days ago

68b08bd3

Older