ROCm/jax

Pull Requests Commits

[ROCm] Fix invalid command in dev build script

mrodden committed 2 years ago

c568f2a6

[ROCm] Docker files cleanup

mrodden committed 2 years ago

4febce00

Merge pull request #24315 from jakevdp:setitem-error

Google-ML-Automation committed 2 years ago

ad99ab1c

Reverts 0788a0a589318b98ae2cb2169af109622ba26cf0

Google-ML-Automation committed 2 years ago

30ba7f37

Merge pull request #24316 from ROCm:ci_build_fix

Google-ML-Automation committed 2 years ago

f0459907

Use Iota order for certain v5e with 8 devices.

Google-ML-Automation committed 2 years ago

5e03a573

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

1df58b18

[Pallas TPU] Fix boolean comparison

ayaka14732 committed 2 years ago

5ac2076f

[Pallas TPU] Support ref reshape.

bythew3i committed 2 years ago

3a7d9137

[ROCm] build script fix

Ruturaj4 committed 2 years ago

a2824862

Improve array setitem error

jakevdp committed 2 years ago

63f4299e

[Mosaic TPU] Support native int4 @ int4

bythew3i committed 2 years ago

a47b7556

Merge pull request #24305 from jakevdp:signbit-doc

Google-ML-Automation committed 2 years ago

94152e8d

Merge pull request #24294 from jakevdp:invert-doc

Google-ML-Automation committed 2 years 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 2 years ago

3a3190fb

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

Google-ML-Automation committed 2 years ago

e461c049

Skip some Shardy-enabled tests if XLA < 292.

belitskiy committed 2 years ago

2f2fd8a3

Improve docs for jnp.invert and related functions

jakevdp committed 2 years ago

dd4a0408

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

Google-ML-Automation committed 2 years ago

2c2c1eeb

[Pallas] Use core_map instead of shard_map for Shmallas

sharadmv committed 2 years ago

cd78c653

Add back missing manual axis to fallback.

pschuh committed 2 years 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 2 years ago

2f6cb89a

Better docs for jnp.signbit

jakevdp committed 2 years ago

a0969868

Internal change.

WindQAQ committed 2 years ago

0788a0a5

Merge pull request #24299 from jakevdp:apply-doc

Google-ML-Automation committed 2 years ago

829e3159

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

Google-ML-Automation committed 2 years ago

dfee4d15

Merge pull request #23467 from abhinavgoel95:patch-4

Google-ML-Automation committed 2 years ago

2a828d5d

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

Tongfei-Guo committed 2 years ago

d621737f

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

ayaka14732 committed 2 years ago

bfc3d3cd

Better docs for jnp.apply_along_axis & apply_over_axes

jakevdp committed 2 years ago

0c307fe7

Older