ROCm/jax

Pull Requests Commits

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

Merge pull request #22616 from rajasekharporeddy:testbranch1

a-googler committed 2 years ago

f32b5d9e

Merge pull request #22555 from rajasekharporeddy:testbranch2

a-googler committed 2 years ago

b8aae435

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

770d5169

[Mosaic GPU] Only split collective TMAs only (multiple) major dimensions

apaszke committed 2 years ago

4f19af91

Merge pull request #22636 from jakevdp:test-typos

a-googler committed 2 years ago

06b199f4

Fix incorrect usage of assertWarns in tests

jakevdp committed 2 years ago

0f136463

Cjkkkk committed 2 years ago

c8b474e1

[Mosaic GPU] Strengthen cluster-related tests by covering more cluster shapes

apaszke committed 2 years ago

dbe8f563

Remove dead code now that xmap is deleted

yashk2810 committed 2 years ago

b6e86c41

lax_numpy_test: fix some typos

jakevdp committed 2 years ago

b0cd3b3e

Delete `xmap` and the `jax.experimental.maps` module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

yashk2810 committed 2 years ago

0d5dae09

Don't broadcast scalar conditions in the jnp.where implementation().

hawkinsp committed 2 years ago

75271016

Skip pallas/gpu_attention_test.py on TPU.

belitskiy committed 2 years ago

d9a7cb44

Simplify implementation of _broadcast_to.

hawkinsp committed 2 years ago

34ce9f21

Fix error message in jax.nn.dot_product_attention when the inputs have different dtypes.

mkazmier committed 2 years ago

61374c92

Reverts 0e17d26b6d81a6b281f55bdd81a6f0ab45efeafe

gnecula committed 2 years ago

0d058ce8

Improved docs for jnp.triu_indices and tril_indices

rajasekharporeddy committed 2 years ago

c8f30875

Merge pull request #22552 from gnecula:pallas_grid

a-googler committed 2 years ago

0e17d26b

[pallas] Simplify handling of BlockMapping and GridMapping

gnecula committed 2 years ago

c5871331

[Mosaic GPU] Move barrier allocation to SMEM scratch specs

apaszke committed 2 years ago

e52dc7ed

Port Triangular Solve to XLA's FFI

Paweł Paruzel committed 2 years ago

54fe6e68

[Mosaic GPU] Allow unions nested inside the smem ref tree

apaszke committed 2 years ago

832eb2d8

[Mosaic GPU] Add sin/cos + unify support for approximate transcendental functions

apaszke committed 2 years ago

6bc79293

[Pallas/TPU] Allow reading DMA semaphores in Pallas

sharadmv committed 2 years ago

cfd9d8f5

Merge pull request #22610 from mattjj:12719

a-googler committed 2 years ago

50c56136

fix memory leak in cond jaxpr tracig

mattjj committed 2 years ago

8db862c0

[XLA:Mosaic] Pass rewrite ctx of apply-vector-layout pass to relayout function.

bythew3i committed 2 years ago

59e944da

Newer Older