ROCm/jax

Pull Requests Commits

Update rules_ml_toolchain reference

alekstheod committed 116 days ago

aaf48382

Move PartitionSpec back to python and intern it.

hawkinsp committed 116 days ago

72b90f40

[Pallas/Mosaic GPU] Implement logic that allows commuting `ReshapeTransform`s with `UntilingTransform`s.

bchetioui committed 116 days ago

4cda3a9d

[Mosaic GPU] Require explicit indexing for barriers initialised with num_barriers > 1.

allanrenucci committed 116 days ago

15371b08

[Mosaic GPU] Disable XLA autotuning for `matmul_test`.

bchetioui committed 116 days ago

3345b109

[pallas:sc] Added `plsc.Indices` which allow specifying the ignored offset value for gathers/scatters

superbobry committed 116 days ago

a7ac409f

[Mosaic GPU] Allow collapsing unit dimensions when issuing TMA loads.

bchetioui committed 116 days ago

a56a2210

[CI] Fix portserver not being shut down on the doc render Sphinx command failing.

belitskiy committed 116 days ago

59e12d5e

Fix cute dsl documentation rendering.

belitskiy committed 116 days ago

df4095e8

Merge pull request #37225 from superbobry:maint

Google-ML-Automation committed 116 days ago

ba9e6746

Merge pull request #37224 from superbobry:main

Google-ML-Automation committed 116 days ago

1daaa13b

[Mosaic GPU] Bump the number of test shards to prevent sporadic timeouts

apaszke committed 116 days ago

cae7baf6

Fixed jupytext failures at head

superbobry committed 116 days ago

1bb0f154

Bumped Pyrefly to 0.63.0

superbobry committed 116 days ago

72434a91

[mgpu] Use `s2f6` for i4->bf16 cast.

chr1sj0nes committed 116 days ago

dfb0ee01

[Mosaic GPU] Always use a transformed input ref in mgpu.vector_load lowering

apaszke committed 116 days ago

be903d99

Pin curl to a specific version to ensure patch applies correctly

meteorcloudy committed 116 days ago

9bbd6ca6

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

Google-ML-Automation committed 116 days ago

7a72f18d

[Mosaic] Add a non-static version of VectorLayout::getImplicitTiledVals

tlongeri committed 116 days ago

04a887cd

Merge pull request #37221 from gnecula:fix_repro

Google-ML-Automation committed 116 days ago

aadcfb1e

Small fix in the hooks for repro generation.

gnecula committed 116 days ago

39988f3d

[Pallas] Have pl.kernel be implemented via mpmd_map not core_map

sharadmv committed 116 days ago

18597c60

[pallas:mgpu] Allow multi-dimensional barriers.

chr1sj0nes committed 116 days ago

d0d53dcb

[Pallas] Add support for closed over Refs and scalars in mpmd map

sharadmv committed 117 days ago

fc06436f

Rename AutoLayout class to AutoLayoutSingleton for some upcoming experiments

yashk2810 committed 117 days ago

249787b3

[Mosaic] In TiledLayoutAttr, add a function to get the number of trailing dims with contiguous tiles

tlongeri committed 117 days ago

16f5b3ee

Merge pull request #36496 from katjasrz:kernel_authoring_documentation

Google-ML-Automation committed 117 days ago

f3258ac9

Update NCCL version 2.28.9 -> 2.29.7.

PatriosTheGreat committed 117 days ago

162c98f4

[Pallas] Support take_along_axis in SC

tlongeri committed 117 days ago

e60c3e9e

Merge pull request #35415 from dlwh:fix-s3-normpath-regex

Google-ML-Automation committed 117 days ago

3e9056a0

Older