ROCm/jax

Pull Requests Commits

Update rules_ml_toolchain as well as xla refs

alekstheod committed 42 days ago

dc589005

[pallas:mgpu] Slightly reworded the docstring of `copy_gmem_to_smem`

superbobry committed 42 days ago

4e35adb4

[pallas:mgpu] Removed the duplicate changelog header

superbobry committed 42 days ago

0e77bca0

Finalize deprecations of several long-deprecated APIs in `jax.interpreters.pxla`.

Jake VanderPlas committed 42 days ago

fbd3214b

[Mosaic TPU] Strengthen verification for TiledLayoutAttr

apaszke committed 42 days ago

b796bc70

[Mosaic GPU] Allow varying the mapping of warps between M and N MMA dims

apaszke committed 42 days ago

c07445d9

[Mosaic GPU] Implement the fully general txmatrix transfer planning

apaszke committed 42 days ago

f4ccaa46

[Mosaic GPU] Use unrealized conversion cast operands directly in dialect_lowering.

golechwierowicz committed 42 days ago

de71c540

[Mosaic GPU][NFC] Refactor in preparation for a more general txmatrix planner

apaszke committed 42 days ago

155406bf

[Mosaic GPU] Improve the txmatrix factorization scheme

apaszke committed 42 days ago

2e07c295

[Mosaic GPU] Add support for ldmatrix/stmatrix with num=4

apaszke committed 42 days ago

a2ceab21

[Mosaic GPU] Add basic support for ldmatrix/stmatrix

apaszke committed 42 days ago

dae2e8e7

[Mosaic GPU] Add support for barriers on Ampere GPUs

apaszke committed 42 days ago

ea11dbfc

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

Google-ML-Automation committed 42 days ago

87602125

[Mosaic GPU] Support PDL griddepcontrol lowering in Warp core_map

rajaanand-g committed 43 days ago

b560a81a

Fix test_sin_to_cos_platform_filtering that's broken on v4i

yashk2810 committed 43 days ago

7f413fba

Merge pull request #39048 from skye-ant:xla_transform_test

Google-ML-Automation committed 43 days ago

ff4265a2

Make xla_transform_test.py run on TPU

skye-ant committed 43 days ago

Verified b1652163

Merge pull request #39020 from jax-ml:merge-jaxpr-closedjaxpr-follow-up

Google-ML-Automation committed 43 days ago

0db0a8b7

Merge pull request #39040 from mattjj:push-nxmtpktpkmyo

Google-ML-Automation committed 43 days ago

6d9bf888

Make `@remat` work nicely with jaxtyping when `f is None`

yashk2810 committed 43 days ago

ddcbdd94

Follow-up to Jaxpr/ClosedJaxpr merge

dougalm committed 43 days ago

cda7a2fb

Add `batch_dim_rule` to CheckpointName and PrimalLeftTangentRight

yashk2810 committed 43 days ago

50734573

finalize deprecation of `jax.core.abstract_token`

Jake VanderPlas committed 43 days ago

67740f3a

Guard the new cpp shard_args under a config until it works for everything.

pschuh committed 43 days ago

acde841d

[hijax] error on specified jit in/out_shardings for hi types

mattjj committed 43 days ago

0062efbf

[Mosaic] Switch from diagnostic to opt<string> for recoverable verification failures to improve compile time

naummo committed 43 days ago

a10d53c0

rollback of rollback of ClosedJaxpr/Jaxpr merge

dougalm committed 43 days ago

10ae8c8b

Finalize deprecation of several `jax.core` APIs

Jake VanderPlas committed 43 days ago

ab087271

[Roll forward] Strip Mosaic kernel debug info and exclude autotune cache dir from cache keys

yashk2810 committed 43 days ago

ad707df0

Older