ROCm/jax

Pull Requests Commits

Reduce number of jobs for jax builds under rocm rbe

alekstheod committed 1 day ago

7c9bf9ef

Convert assert into an Exception

yashk2810 committed 2 days ago

079f77e8

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

Google-ML-Automation committed 2 days ago

dbda9168

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

Google-ML-Automation committed 2 days ago

192aa201

[fuser] Add support for tile with squeezed block dims

brianwa84 committed 2 days ago

083cea77

Fix pjit warning stack depth

brianwa84 committed 2 days ago

66666620

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

cperivol committed 2 days ago

7f5daab0

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

allanrenucci committed 2 days ago

8010a97f

Fix test failure in newly added test.

hawkinsp committed 2 days ago

44be67a6

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

bchetioui committed 2 days ago

417f760d

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

bchetioui committed 2 days ago

2842ecb5

[mosaic] Add dynamic_sizes support to tpu.reinterpret_cast

superbobry committed 2 days ago

a0d447f5

[mosaic_gpu] Added layout inference for the synchronous MMA op

superbobry committed 2 days ago

64b7a2a9

[doc] fix some inconsistencies around partial and jit

jakevdp committed 2 days ago

a8fea649

[mosaic_gpu] Added an op for synchronous Ampere MMA

superbobry committed 2 days ago

bf077adf

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

golechwierowicz committed 2 days ago

7b45bc03

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

Google-ML-Automation committed 2 days ago

b1e3e7b6

Reverts c3da028b6bd1e791c44e1d9d36b5771e1304cb46

junwhanahn committed 3 days ago

8e4f9e62

[bazel] visibility modifications

Jake VanderPlas committed 3 days ago

33a278b8

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

Google-ML-Automation committed 3 days ago

ba348c6c

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

yashk2810 committed 3 days ago

77944e22

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

Google-ML-Automation committed 3 days ago

862f2025

use empty not zeros in sum type representation

mattjj committed 3 days ago

15fca029

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

brills committed 3 days ago

f65eb481

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

yashk2810 committed 3 days ago

c3da028b

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

hawkinsp committed 3 days ago

46ea6d33

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

Google-ML-Automation committed 3 days ago

3e00b7cc

[pallas:mgpu] Fix flaky test_grid_dep_control

Google-ML-Automation committed 3 days ago

ecb0f7dd

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

hawkinsp committed 3 days ago

7e9ead37

Support tiling and swizzle transforms in GPU Interpret Mode.

jburnim committed 3 days ago

a4d60cd1

Older