ROCm/jax

Pull Requests Commits

[ROCm]: Add support to continue on fail, fix script paths and update Dockerfile to add necessary packages

Rahul Batra committed 2 years ago

7d6fa3c0

[mosaic_gpu] arrive_expect_tx() also accepts index typed values

cperivol committed 2 years ago

d54bf529

Increase shard count for //third_party/py/jax/tests:lax_numpy_ufuncs_test_cpu.

belitskiy committed 2 years ago

ba50e774

Fixed Pallas Mosaic GPU tests

superbobry committed 2 years ago

99a8b92a

[Mosaic GPU] Add support for clusters in the matmul example

apaszke committed 2 years ago

2e6da35e

[Mosaic GPU] Simplify the matmul example

apaszke committed 2 years ago

e59303cf

[Mosaic GPU] Minor cleanups in the matmul example

apaszke committed 2 years ago

be9cc807

#sdy Support with_sharding_constraint lowering through Shardy.

bartchr808 committed 2 years ago

b00f978f

[Pallas/TPU] Fix bug with LocalMask grid shrinking

Google-ML-Automation committed 2 years ago

f15f9717

Merge pull request #22642 from ROCm:ci_jax_exp

a-googler committed 2 years ago

e14752c0

Merge pull request #22605 from Cjkkkk:add_sm86_sm89_flash_attention

a-googler committed 2 years ago

2dadbd7e

[pallas] Stop exporting jax.experimental.pallas.pallas

gnecula committed 2 years ago

e6066e52

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

gnecula committed 2 years ago

4063373b

remove inlined jax.nn.initializers definitions, resolving TODO of levskaya et al

mattjj committed 2 years ago

c8ea86c9

Merge pull request #22628 from hawkinsp:broadcast2

a-googler committed 2 years ago

76b4c70c

Simplify pjit's batching rule now that xmap is deleted. Also do cleanup around adding manual axes under shard_map

yashk2810 committed 2 years ago

51e27923

Merge pull request #21069 from mattjj:remove-named-shapes

a-googler committed 2 years ago

086b500d

remove named_shapes (since xmap is now gone)

mattjj committed 2 years ago

3f9eb404

[pallas] Use the same primitive `run_scoped_p` for moth mosaic and mosaic_gpu

cperivol committed 2 years ago

80a193d5

[MosaicGPU] Expose the BarrierRef type so we can type our variables outside of mosaic.

cperivol committed 2 years ago

aa991131

Merge pull request #22634 from jakevdp:array-api-cleanup

a-googler committed 2 years ago

b5010a2d

[Mosaic TPU] (8,128),-2 -> (8,128) for non-zero and replicated 2nd minor offset

tlongeri committed 2 years ago

220ec2aa

Simplify promote_shapes.

hawkinsp committed 2 years ago

52fa165d

Merge pull request #22625 from hawkinsp:broadcast

a-googler committed 2 years ago

f1cfd99f

[array API] clean up some unused/unnecessary code

jakevdp committed 2 years ago

b8bb869b

Merge pull request #22640 from jakevdp:fix-assert-warns

a-googler committed 2 years ago

422e033d

Merge pull request #22629 from hawkinsp:broadcast3

a-googler committed 2 years ago

02235eb5

Add `canonical` arg to Rotation.as_quat() and switch .inv() to use the quaternion conjugate.

SiegeLordEx committed 2 years ago

5ba26953

Update shard_map(jit) to properly set manual_axes on in_shardings and out_shardings of the nested jit. This avoids a problem where the jit returns {manaual} and then this gets passed to ShardToFull (manual is already considered a full sharding).

pschuh committed 2 years ago

be6b77cc

[ROCM] rocm plugin is no longer experiemental!

Ruturaj4 committed 2 years ago

f7039b91

Older