ROCm/jax

Pull Requests Commits

Reduce number of jobs for jax builds under rocm rbe

alekstheod committed 71 days ago

7c9bf9ef

Convert assert into an Exception

yashk2810 committed 72 days ago

079f77e8

Merge pull request #38295 from jakevdp:doc-jit-partial

Google-ML-Automation committed 72 days ago

dbda9168

[Pallas][GPU kernel interpreter] Properly support `collective_axes` for allocations made by `run_scoped`.

Google-ML-Automation committed 72 days ago

192aa201

[fuser] Add support for tile with squeezed block dims

brianwa84 committed 72 days ago

083cea77

Fix pjit warning stack depth

brianwa84 committed 72 days ago

66666620

[pallas:mgpu] Expose TMA_INDICES_4 to pallas and fix layout inference for it.

cperivol committed 72 days ago

7f5daab0

[Mosaic GPU] Enable column slicing for TMEM layouts with varying row dimensions.

allanrenucci committed 72 days ago

8010a97f

Fix test failure in newly added test.

hawkinsp committed 72 days ago

44be67a6

[Pallas/Mosaic GPU][NFC] Remove unnecessary kernel name from `test_griddepcontrol`.

bchetioui committed 72 days ago

417f760d

[Pallas/Mosaic GPU][NFC] Clean up `test_griddepcontrol`.

bchetioui committed 72 days ago

2842ecb5

[mosaic] Add dynamic_sizes support to tpu.reinterpret_cast

superbobry committed 72 days ago

a0d447f5

[mosaic_gpu] Added layout inference for the synchronous MMA op

superbobry committed 72 days ago

64b7a2a9

[doc] fix some inconsistencies around partial and jit

jakevdp committed 72 days ago

a8fea649

[mosaic_gpu] Added an op for synchronous Ampere MMA

superbobry committed 72 days ago

bf077adf

[Mosaic GPU] Safely construct certain layouts instead of failing fast with exception.

golechwierowicz committed 72 days ago

7b45bc03

Update XLA dependency to use revision http://github.com/openxla/xla/commit/26880785d1a0b68d1f08398117f0d557bb656430

Google-ML-Automation committed 72 days ago

b1e3e7b6

Reverts c3da028b6bd1e791c44e1d9d36b5771e1304cb46

junwhanahn committed 72 days ago

8e4f9e62

[bazel] visibility modifications

Jake VanderPlas committed 72 days ago

33a278b8

Merge pull request #38279 from mattjj:empty-not-zeros-in-sum-type-representation

Google-ML-Automation committed 72 days ago

ba348c6c

Improve error message for sharding rules when the dimension does not divide mesh axes it's sharded on

yashk2810 committed 72 days ago

77944e22

Merge pull request #38168 from mattjj:accum-ct-avals

Google-ML-Automation committed 72 days ago

862f2025

use empty not zeros in sum type representation

mattjj committed 72 days ago

15fca029

Bug fix: do not create real (non-symbolic) zero arrays for tangents.

brills committed 72 days ago

f65eb481

Use `lax.empty` in `_empty_array` and add `out_memory_space` to `lax.empty`.

yashk2810 committed 72 days ago

c3da028b

[pallas] When unrolling loops during mosaic lowering, emit the loop body out of line and repeatedly inline it.

hawkinsp committed 72 days ago

46ea6d33

[Pallas:MGPU] Allow dynamic shapes in nd_loop and planar_snake.

Google-ML-Automation committed 72 days ago

3e00b7cc

[pallas:mgpu] Fix flaky test_grid_dep_control

Google-ML-Automation committed 72 days ago

ecb0f7dd

In Primitive.bind, canonicalize first before calling typeof and fuse the loops that compute them.

hawkinsp committed 72 days ago

7e9ead37

Support tiling and swizzle transforms in GPU Interpret Mode.

jburnim committed 72 days ago

a4d60cd1

Older