ROCm/jax

Pull Requests Commits

Use CloudFront download for ROCm plugin/PJRT wheels instead of S3

psanal35 committed 79 days ago

444bb1ea

[NFC] Clean up a bit of logic in take_along_axis

tlongeri committed 79 days ago

798f4884

Merge pull request #36772 from superbobry:pyrefly-2

Google-ML-Automation committed 79 days ago

3c4fcfe3

[Mosaic GPU] Fix bug in the implementation of `tile_strides` in the presence of duplicate strides.

bchetioui committed 79 days ago

94748a15

Change deserialize_portable_artifacts to return an MLIR Module directly rather than forming a string.

hawkinsp committed 79 days ago

af938de7

[XLA] Make sure that executable serialization is deterministic.

hawkinsp committed 79 days ago

07d5c405

[Pallas/Mosaic GPU] Fix bug in computation of profiler buffer allocation size.

bchetioui committed 79 days ago

549a3e37

[Mosaic GPU][NFC] Fold the computation of iteration spaces over tiles inside `plan_tiled_transfer`.

bchetioui committed 79 days ago

c7a22c4a

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

PatriosTheGreat committed 79 days ago

71b0e8d4

Upgraded to Pyrefly 0.61.0

superbobry committed 79 days ago

613d58b5

[consts] Fix const_lowering

gnecula committed 79 days ago

75e3f307

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

golechwierowicz committed 79 days ago

60e44d27

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

chr1sj0nes committed 79 days ago

c5dd0070

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

bchetioui committed 79 days ago

c86c9c2a

Merge pull request #36758 from superbobry:pyrefly-2

Google-ML-Automation committed 79 days ago

cbc2d32f

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

golechwierowicz committed 79 days ago

b9f5a7ef

Removed a few stray `pytype: enable` directives

superbobry committed 79 days ago

c364a892

[Mosaic GPU] Add MultimemLoadReduceOp to the dialect.

golechwierowicz committed 79 days ago

d1a0bd20

[Mosaic GPU] Add support for async cluster atomics

apaszke committed 79 days ago

e3af3e89

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

Google-ML-Automation committed 79 days ago

b160451a

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

superbobry committed 79 days ago

c79a875f

[Mosaic] Fix a typo in the DMA verification error

naummo committed 79 days ago

761d5eea

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

Google-ML-Automation committed 79 days ago

6165fcf4

Allow passing an aval directly to TypedNdArray.

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

8f39e8fc

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

Rifur13 committed 80 days ago

53f6c3a5

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

Google-ML-Automation committed 80 days ago

788ef131

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

C-you-know committed 80 days ago

f3e5b9af

Concatenate caller and callee tracebacks in jaxpr equation profiles.

hawkinsp committed 80 days ago

4e442000

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

belitskiy committed 80 days ago

2ff3d741

Older