ROCm/jax

Pull Requests Commits

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

Google-ML-Automation committed 7 days ago

4a592c9d

[ci] silence pyrefly error at head

jakevdp committed 7 days ago

d0ce0102

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

Google-ML-Automation committed 8 days ago

fc393eca

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

Google-ML-Automation committed 9 days ago

7b6360dd

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

tlongeri committed 9 days ago

23bede0c

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

tlongeri committed 9 days ago

d1495d28

Reverts a1ff74b306dc37e979bee7c8b56e1d9eb425bcce

gnecula committed 9 days ago

cfb68889

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

tlongeri committed 9 days ago

2130ed7f

Use `ct_zero_from_primal` in shard_map_transpose to create zeros from residuals

yashk2810 committed 9 days ago

82f4db2a

Reverts c43e68919106c28b379ec9353439809e2ce7f797

gnecula committed 9 days ago

ab202431

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

tlongeri committed 9 days ago

dedbc60d

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

ashishenoyp committed 9 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 9 days ago

acbca16b

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

Google-ML-Automation committed 9 days ago

e578b6d4

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

Google-ML-Automation committed 9 days ago

1808f592

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

tlongeri committed 9 days ago

27650dfb

Bump pydata-sphinx-theme from 0.14.4 to 0.16.1

dependabot[bot] committed 9 days ago

Verified 2fee16fa

Merge pull request #35689 from jakevdp:pyrefly-pallas

Google-ML-Automation committed 9 days ago

b69238e7

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

alekstheod committed 9 days ago

9a0eddda

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

ybaturina committed 9 days ago

80f3cf17

[pyrefly] type check jax/experimental/pallas

jakevdp committed 9 days ago

323ec0b8

Fix flaky test in ffi_test.

hawkinsp committed 9 days ago

6f3c7794

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

belitskiy committed 9 days ago

00b08ccb

Update HERMETIC_CUDA_UMD_VERSION 13.0.0 -> 13.0.2.

ybaturina committed 9 days ago

c80435cb

Merge pull request #35677 from gnecula:ser_abs_device

Google-ML-Automation committed 9 days ago

c43e6891

[export] Add support for serializating AbstractMesh.abstract_device

gnecula committed 9 days ago

5498417d

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

apaszke committed 9 days ago

b7e106f4

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

bchetioui committed 9 days ago

15a6469e

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

bchetioui committed 9 days ago

8f232c94

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

allanrenucci committed 9 days ago

68b08bd3

Older