ROCm/jax

Pull Requests Commits

Remove wheel upload for easier testing.

tsrw2048 committed 151 days ago

0e097bd6

Remove non-ROCm workflows for easier testing.

tsrw2048 committed 151 days ago

6505bc30

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

tsrw2048 committed 151 days ago

937e7998

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

PatriosTheGreat committed 151 days ago

71b0e8d4

[consts] Fix const_lowering

gnecula committed 151 days ago

75e3f307

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

golechwierowicz committed 151 days ago

60e44d27

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

chr1sj0nes committed 151 days ago

c5dd0070

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

bchetioui committed 151 days ago

c86c9c2a

Merge pull request #36758 from superbobry:pyrefly-2

Google-ML-Automation committed 151 days ago

cbc2d32f

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

golechwierowicz committed 151 days ago

b9f5a7ef

Removed a few stray `pytype: enable` directives

superbobry committed 151 days ago

c364a892

[Mosaic GPU] Add MultimemLoadReduceOp to the dialect.

golechwierowicz committed 151 days ago

d1a0bd20

[Mosaic GPU] Add support for async cluster atomics

apaszke committed 151 days ago

e3af3e89

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

Google-ML-Automation committed 151 days ago

b160451a

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

superbobry committed 151 days ago

c79a875f

[Mosaic] Fix a typo in the DMA verification error

naummo committed 151 days ago

761d5eea

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

Google-ML-Automation committed 151 days ago

6165fcf4

Allow passing an aval directly to TypedNdArray.

hawkinsp committed 152 days 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 152 days ago

8f39e8fc

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

Rifur13 committed 152 days ago

53f6c3a5

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

Google-ML-Automation committed 152 days ago

788ef131

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

C-you-know committed 152 days ago

f3e5b9af

Concatenate caller and callee tracebacks in jaxpr equation profiles.

hawkinsp committed 152 days ago

4e442000

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

belitskiy committed 152 days ago

2ff3d741

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

Google-ML-Automation committed 152 days ago

a4c77b21

Merge pull request #36712 from superbobry:pyrefly-2

Google-ML-Automation committed 152 days ago

6f2498b4

Removed outdated comment

superbobry committed 152 days ago

a7e3db8a

Fix grad-of-conv_general_dilated with preferred_element_type

jakevdp committed 152 days ago

585f9f2b

Removed many blanket `type: ignore` directives

superbobry committed 152 days ago

40cf5e09

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

hawkinsp committed 152 days ago

2eba8d26

Older