ROCm/jax

Pull Requests Commits

Remove wheel upload for easier testing.

tsrw2048 committed 7 hours ago

0e097bd6

Remove non-ROCm workflows for easier testing.

tsrw2048 committed 7 hours ago

6505bc30

[ROCm] Add `remote_download_outputs` option at build invocation.

tsrw2048 committed 8 hours ago

937e7998

[Mosaic:GPU] Use separate allocation to register multicast handler, instead of entire memory range.

PatriosTheGreat committed 9 hours ago

71b0e8d4

[consts] Fix const_lowering

gnecula committed 10 hours ago

75e3f307

[Mosaic GPU] Add pallas and dialect WG lowering for multimem_load_reduce.

golechwierowicz committed 10 hours ago

60e44d27

[pallas] Allow `loop` functions to have a loop carry.

chr1sj0nes committed 11 hours ago

c5dd0070

[Pallas/Mosaic GPU] Support warp-specialized `tcgen05_mma` under warpgroup semantics.

bchetioui committed 11 hours ago

c86c9c2a

Merge pull request #36758 from superbobry:pyrefly-2

Google-ML-Automation committed 12 hours ago

cbc2d32f

[Mosaic GPU] Add layout inference rule to the MultimemLoadReduceOp.

golechwierowicz committed 12 hours ago

b9f5a7ef

Removed a few stray `pytype: enable` directives

superbobry committed 13 hours ago

c364a892

[Mosaic GPU] Add MultimemLoadReduceOp to the dialect.

golechwierowicz committed 13 hours ago

d1a0bd20

[Mosaic GPU] Add support for async cluster atomics

apaszke committed 13 hours ago

e3af3e89

Merge pull request #36675 from olupton:fix-execution-on-more-than-2-devices

Google-ML-Automation committed 13 hours ago

b160451a

Removed or revised a few more `type: ignore`s

superbobry committed 13 hours ago

c79a875f

[Mosaic] Fix a typo in the DMA verification error

naummo committed 14 hours ago

761d5eea

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

Google-ML-Automation committed 14 hours ago

6165fcf4

Allow passing an aval directly to TypedNdArray.

hawkinsp committed 20 hours ago

408bd93e

Allow reshape split `f32[4@(x,y)] -> f32[2@x, 2@y]` which is a no-comms operation. Similarly allow merge `f32[2@x, 2@y] -> f32[4@(x,y)]`.

yashk2810 committed 22 hours ago

8f39e8fc

[Pallas MGPU] We currently make no promises about the out of bounds values of GMEM->SMEM copies.

Rifur13 committed 22 hours ago

53f6c3a5

Merge pull request #34084 from C-you-know:feature-weighted-quantile

Google-ML-Automation committed 23 hours ago

788ef131

feat: Add weights parameter to quantile and percentile functions (#32647)

C-you-know committed 1 day ago

f3e5b9af

Concatenate caller and callee tracebacks in jaxpr equation profiles.

hawkinsp committed 1 day ago

4e442000

Remove the old nightly pin for LibTPU since the segfault regression is gone.

belitskiy committed 1 day ago

2ff3d741

Merge pull request #36734 from jakevdp:conv-general-dilated-grad

Google-ML-Automation committed 1 day ago

a4c77b21

Merge pull request #36712 from superbobry:pyrefly-2

Google-ML-Automation committed 1 day ago

6f2498b4

Removed outdated comment

superbobry committed 1 day ago

a7e3db8a

Fix grad-of-conv_general_dilated with preferred_element_type

jakevdp committed 1 day ago

585f9f2b

Removed many blanket `type: ignore` directives

superbobry committed 1 day ago

40cf5e09

Don't dump the MLIR module unless we have debug logging enabled in export.py

hawkinsp committed 1 day ago

2eba8d26

Older