ROCm/jax

Pull Requests Commits

Merge pull request #25296 from hawkinsp:postrelease

Google-ML-Automation committed 1 year ago

45159a2f

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1ca8903a

Merge branch 'release/0.4.36' after release

hawkinsp committed 1 year ago

ab02bf87

[pallas] fix jumble test flakiness

Google-ML-Automation committed 1 year ago

84f3f992

[Mosaic:TPU] Fix elementwise inference with i1s

tlongeri committed 1 year ago

651ab188

Merge pull request #25269 from justinjfu:pallas_docs_cleanup

Google-ML-Automation committed 1 year ago

d782b246

[Pallas] Fix shard_axis in dma_start interpret mode rule.

justinjfu committed 1 year ago

fd42b561

[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

[shape_poly] Fix the handling of __pow__ for symbolic dimensions

gnecula committed 1 year ago

4e17bea9

[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

Newer Older