ROCm/jax

Pull Requests Commits

Fix run_multi_gpu script multi-gpu issue and refactor code

Ruturaj4 committed 346 days ago

1d986764

Fix hybrid type for ROCm

Ruturaj4 committed 360 days ago

afd4ab63

Update documentation and add setup.py pypi bindings

Ruturaj4 committed 360 days ago

7173d445

[ROCm] Remove exit on error (#167)

Ruturaj4 committed 360 days ago

bd1b3329

[ROCm] ci build and dockerfile changes

Ruturaj4 committed 360 days ago

7a7db89f

BUILD file fixes and linter run

Ruturaj4 committed 360 days ago

71e6e066

[Pallas] Fix shard_axis in dma_start interpret mode rule.

justinjfu committed 1 year ago

259194a6

JAX release 0.4.36.

hawkinsp committed 1 year ago

7e6620a5

[Mosaic:TPU] Fix fully replicated relayout

tlongeri committed 1 year ago

23d5c10f

[jax:custom_partitioning] Implement SdyShardingRule to support

bixia1 committed 1 year ago

2a4a0e8d

Merge pull request #25290 from jakevdp:reduction-where

Google-ML-Automation committed 1 year ago

f73fa7a7

Merge pull request #25271 from jakevdp:fix-vector-norm

Google-ML-Automation committed 1 year ago

a71f9a62

[JAX] Add end-to-end execution support in colocated Python API

hyeontaek committed 1 year ago

e20a483b

jnp.linalg.vector_norm: properly support multiple axes

jakevdp committed 1 year ago

aaaee63a

jax.numpy: require boolean dtype for where argument

jakevdp committed 1 year ago

29a8cce6

[export] Removed __gpu$xla.gpu.triton (Pallas GPU) from the list of custom calls with guaranteed compatibility.

gnecula committed 1 year ago

3f5f3e1c

[pallas:mosaic_gpu] Removed unnecessarily strict check in `emit_pipeline`

superbobry committed 1 year ago

4a41aa0a

[shape_poly] Remove some deprecated kwargs

gnecula committed 1 year ago

5fe5206b

[pallas:mosaic_gpu] Do not store the grid mapping in `ModuleContext`

superbobry committed 1 year ago

e5102957

[Mosaic GPU] Always annotate block initialization in the profiles

apaszke committed 1 year ago

d034680f

[Mosaic TPU] Add support for modeling loads/stores and fix minor issues in model extraction

apaszke committed 1 year ago

d5ead570

Reverts 73962b740890a728295fa09f515dcf96cb820822

chr1sj0nes committed 1 year ago

569c2a3c

Merge pull request #25276 from mattjj:remove-vestigial-reducing-transposes

Google-ML-Automation committed 1 year ago

39d73a68

[pallas:mosaic_gpu] Removed leftover debugging code

superbobry committed 1 year ago

03861d43

[Mosaic GPU] Remove expect_wait from Barrier.wait

apaszke committed 1 year ago

c965ffbf

[AutoPGLE] Add multi-process test case

Google-ML-Automation committed 1 year ago

7214a3a8

[Mosaic:TPU] Add relayout for adding minor implicit dim and relax some offset restrictions on similar shape cast

tlongeri committed 1 year ago

8163e74e

remove vestigial ad.reducing_transposes table

mattjj committed 1 year ago

6172a1f1

[Mosaic:TPU] Lift offset restrictions on single-row (1, 128) -> (8, 128) 32-bit replicated retiling

tlongeri committed 1 year ago

10116874

More thorough propagation of host linear layout. Currently linear layout on host

Google-ML-Automation committed 1 year ago

f160df04

Older