ROCm/jax

Pull Requests Commits

Update ROCm plugin to use ffi_registrations() API from main

alekstheod committed 116 days ago

39eb82a0

Fix plugin init script

alekstheod committed 116 days ago

a521096d

Don't create residual zeros at all in fun_trans_callable in shard_map_transpose because we are going to return Zero ct for them anyway so no need to initialize actual zero arrays for them.

yashk2810 committed 117 days ago

67edb448

[pallas:mosaic] De-duplicated the implementation of various `BufferedRef` APIs

superbobry committed 117 days ago

c014b6ef

Merge pull request #35384 from Shreyas-S-809:fix35105

Google-ML-Automation committed 117 days ago

30017a87

[Mosaic GPU] Fix predicate for `tcgen05_mma` lowering.

allanrenucci committed 118 days ago

3f7d9106

Merge pull request #35223 from Harshadev-24:add-hankel-matrix

Google-ML-Automation committed 118 days ago

9334db25

[Mosaic GPU] Add explicit control over the rounding mode to FragmentedArray.astype

apaszke committed 118 days ago

0b846bdb

[ragged-paged-attn] add parsing for tpu7x

Google-ML-Automation committed 118 days ago

def09916

Require arguments to Primitive.bind() to pass typeof().

hawkinsp committed 118 days ago

784a37c2

Merge pull request #35714 from superbobry:pyrefly

Google-ML-Automation committed 118 days ago

cd07abf1

[Mosaic GPU] Minor logic cleanup in FragmentedArray.reduce

apaszke committed 118 days ago

0aa31d19

[Pallas:MGPU] Add support for SMEM atomics

apaszke committed 118 days ago

d512581a

[Mosaic GPU] Implement layout inference for sparse `tcgen05.mma`.

bchetioui committed 118 days ago

91a0fd18

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

Google-ML-Automation committed 118 days ago

d883023f

Show all mismatches in `flatten_axes_resources()` for pytree prefix errors instead of just the first mismatch. This prevents multiple roundtrips just to identify all pytree mismatches one by one.

ashishenoyp committed 118 days ago

50476ee1

[pyrefly] Excluded jax/_src/internal_test_util/export_back_compat_test_data

superbobry committed 118 days ago

cc7c25bb

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

Google-ML-Automation committed 118 days ago

4a592c9d

[ci] silence pyrefly error at head

jakevdp committed 118 days ago

d0ce0102

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

Google-ML-Automation committed 119 days ago

fc393eca

scipy.linalg: add hankel special matrix

Harshadev-24 committed 119 days ago

77d844bd

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

Google-ML-Automation committed 120 days ago

7b6360dd

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

tlongeri committed 120 days ago

23bede0c

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

tlongeri committed 120 days ago

d1495d28

Reverts a1ff74b306dc37e979bee7c8b56e1d9eb425bcce

gnecula committed 120 days ago

cfb68889

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

tlongeri committed 120 days ago

2130ed7f

Use `ct_zero_from_primal` in shard_map_transpose to create zeros from residuals

yashk2810 committed 120 days ago

82f4db2a

Reverts c43e68919106c28b379ec9353439809e2ce7f797

gnecula committed 120 days ago

ab202431

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

tlongeri committed 120 days ago

dedbc60d

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

ashishenoyp committed 120 days ago

caee8a92

Older