ROCm/jax

Pull Requests Commits

[ROCM]: Generating pytest html logs from unit-tests.

zahiqbal committed 2 years ago

ef7694f2

Add an internal `_device_list` parameter to GSPMDSharding so that we can save on the initialization cost of PyDeviceList when creating `GSPMDSharding` from other shardings

yashk2810 committed 2 years ago

a63197fe

Merge pull request #19492 from gnecula:poly_tests

a-googler committed 2 years ago

a74b04a4

[shape_poly] Add more tests for reasoning about inequalities.

gnecula committed 2 years ago

0bd511d6

Update XLA dependency to use revision

a-googler committed 2 years ago

eab0dd19

Preserve single device NamedSharding/PositionalSharding on the output instead of always return SingleDeviceShardings.

yashk2810 committed 2 years ago

6f96c963

Merge pull request #19145 from 8bitmp3:jax-docs-jax-primitives

a-googler committed 2 years ago

eb59716e

Upgrade JAX primitives doc

8bitmp3 committed 2 years ago

a295b785

Merge pull request #19489 from skye:version

a-googler committed 2 years ago

f67d4823

Update version numbers for 0.4.24 release

skye committed 2 years ago

368f594c

[xla:gpu] Make cu_threefry2x32 custom call compatible with command buffers

anlunx committed 2 years ago

88f5eaca

Merge pull request #19484 from jakevdp:full-tests

a-googler committed 2 years ago

85fc4a95

[Pallas] Add support for trivial vmap of scalar prefetch

sharadmv committed 2 years ago

037bc5ed

Test: move lax.full sharding tests to multi_device_test

jakevdp committed 2 years ago

8476b44a

Merge pull request #19481 from gnecula:poly_changelog

a-googler committed 2 years ago

7e80aac3

[shape_poly] Forgot to update CHANGELOG for #19235.

gnecula committed 2 years ago

af306680

Merge pull request #19473 from jakevdp:shard-map-run

a-googler committed 2 years ago

c2863726

Merge pull request #19285 from gnecula:poly_constraints

a-googler committed 2 years ago

8708d009

The Triton autotuner ignores configs that use too much shmem

a-googler committed 2 years ago

5761e393

[shape_poly] Add support for symbolic constraints on dimension variables

gnecula committed 2 years ago

24201ef9

Exclude modules with host callbacks from inter-host sharing.

a-googler committed 2 years ago

34d22fc4

Update XLA dependency to use revision

a-googler committed 2 years ago

f4b483cf

Merge pull request #19474 from mattjj:shmap-disable-jit-fix

a-googler committed 2 years ago

507239d0

[shard-map] fix disable_jit-of-shmap

mattjj committed 2 years ago

5275ef6e

DOC: run shard_map docs as part of the doc build

jakevdp committed 2 years ago

a3abd178

Merge pull request #19466 from jakevdp:full-like-sharding

a-googler committed 2 years ago

70aec849

Merge pull request #19472 from mattjj:shmap-tutorial-typos

a-googler committed 2 years ago

f21022b1

Merge pull request #19438 from sharadmv:add-init

a-googler committed 2 years ago

1df92c95

[shard-map] fix typos in docs

mattjj committed 2 years ago

a61f7807

[Pallas/TPU] Add vmem limit test

sharadmv committed 2 years ago

76554bc9

Older