ROCm/jax

Pull Requests Commits

Improve docs for lax.gather & lax.scatter

jakevdp committed 1 year ago

284ca8bc

tile_docstring_added

selamw1 committed 1 year ago

24b6f509

Merge pull request #24315 from jakevdp:setitem-error

Google-ML-Automation committed 1 year ago

ad99ab1c

[Mosaic GPU] Skip OSS Flash Attention test unless running on sm90

andportnoy committed 1 year ago

0da52cd1

Reverts 0788a0a589318b98ae2cb2169af109622ba26cf0

Google-ML-Automation committed 1 year ago

30ba7f37

Merge pull request #24316 from ROCm:ci_build_fix

Google-ML-Automation committed 1 year ago

f0459907

Use Iota order for certain v5e with 8 devices.

Google-ML-Automation committed 1 year ago

5e03a573

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1df58b18

[Pallas TPU] Fix boolean comparison

ayaka14732 committed 1 year ago

5ac2076f

[Pallas TPU] Support ref reshape.

bythew3i committed 1 year ago

3a7d9137

[ROCm] build script fix

Ruturaj4 committed 1 year ago

a2824862

Improve array setitem error

jakevdp committed 1 year ago

63f4299e

[Mosaic TPU] Support native int4 @ int4

bythew3i committed 1 year ago

a47b7556

Merge pull request #24305 from jakevdp:signbit-doc

Google-ML-Automation committed 1 year ago

94152e8d

Merge pull request #24294 from jakevdp:invert-doc

Google-ML-Automation committed 1 year ago

87d8f381

Fix typo in Pallas TPU matmul doc. I think the logical layout of the input array is non-transposed, rather than transposed?

18praveenb committed 1 year ago

3a3190fb

Merge pull request #23684 from simonster:sjk/fix-prefix-error

Google-ML-Automation committed 1 year ago

e461c049

Skip some Shardy-enabled tests if XLA < 292.

belitskiy committed 1 year ago

2f2fd8a3

Improve docs for jnp.invert and related functions

jakevdp committed 1 year ago

dd4a0408

Merge pull request #24251 from dfm:dot-algorithm-jax2tf

Google-ML-Automation committed 1 year ago

2c2c1eeb

[Pallas] Use core_map instead of shard_map for Shmallas

sharadmv committed 1 year ago

cd78c653

Add back missing manual axis to fallback.

pschuh committed 1 year ago

b0768906

Add a private property to NamedSharding called `_logical_device_ids` which allows you to pass a custom `tile_assignment_devices()` equivalent.

yashk2810 committed 1 year ago

2f6cb89a

Better docs for jnp.signbit

jakevdp committed 1 year ago

a0969868

Internal change.

WindQAQ committed 1 year ago

0788a0a5

Merge pull request #24299 from jakevdp:apply-doc

Google-ML-Automation committed 1 year ago

829e3159

Merge pull request #24258 from dfm:remove-jaxlib-version-checks

Google-ML-Automation committed 1 year ago

dfee4d15

Merge pull request #23467 from abhinavgoel95:patch-4

Google-ML-Automation committed 1 year ago

2a828d5d

[XLA:Collective] Expose a factory for constructing HLOSharding with explicit device ordering.

Tongfei-Guo committed 1 year ago

d621737f

[Pallas TPU] Add lowerings for scalar `sin`, `cos`, `tan` and `tanh`

ayaka14732 committed 1 year ago

bfc3d3cd

Newer Older