ROCm/jax

Pull Requests Commits

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

apaszke committed 1 year ago

2e6da35e

[Mosaic GPU] Simplify the matmul example

apaszke committed 1 year ago

e59303cf

[Mosaic GPU] Minor cleanups in the matmul example

apaszke committed 1 year ago

be9cc807

#sdy Support with_sharding_constraint lowering through Shardy.

bartchr808 committed 1 year ago

b00f978f

[Pallas/TPU] Fix bug with LocalMask grid shrinking

Google-ML-Automation committed 1 year ago

f15f9717

Fix jnp.triu_indices_from and jnp.tril_indices_from to emulate NumPy's behavior for arrays other than 2-D

rajasekharporeddy committed 1 year ago

4ff69d38

Merge pull request #22642 from ROCm:ci_jax_exp

a-googler committed 1 year ago

e14752c0

Merge pull request #22605 from Cjkkkk:add_sm86_sm89_flash_attention

a-googler committed 1 year ago

2dadbd7e

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

gnecula committed 1 year ago

e6066e52

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

gnecula committed 1 year ago

4063373b

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

mattjj committed 1 year ago

c8ea86c9

Merge pull request #22628 from hawkinsp:broadcast2

a-googler committed 1 year ago

76b4c70c

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

yashk2810 committed 1 year ago

51e27923

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

a-googler committed 1 year ago

086b500d

remove named_shapes (since xmap is now gone)

mattjj committed 1 year ago

3f9eb404

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

cperivol committed 1 year ago

80a193d5

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

cperivol committed 1 year ago

aa991131

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

a-googler committed 1 year ago

b5010a2d

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

tlongeri committed 1 year ago

220ec2aa

Simplify promote_shapes.

hawkinsp committed 1 year ago

52fa165d

Merge pull request #22625 from hawkinsp:broadcast

a-googler committed 1 year ago

f1cfd99f

[array API] clean up some unused/unnecessary code

jakevdp committed 1 year ago

b8bb869b

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

a-googler committed 1 year ago

422e033d

Merge pull request #22629 from hawkinsp:broadcast3

a-googler committed 1 year ago

02235eb5

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

SiegeLordEx committed 1 year 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 1 year ago

be6b77cc

[ROCM] rocm plugin is no longer experiemental!

Ruturaj4 committed 1 year ago

f7039b91

Merge pull request #22616 from rajasekharporeddy:testbranch1

a-googler committed 1 year ago

f32b5d9e

Merge pull request #22555 from rajasekharporeddy:testbranch2

a-googler committed 1 year ago

b8aae435

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

770d5169

Newer Older