ROCm/jax

Pull Requests Commits

[ROCm] apt update

Ruturaj4 committed 1 year ago

937d79e3

[sharding_in_types] Add `dot_general` sharding rule. We only handle the simple cases and rely on xla to insert the collectives.

yashk2810 committed 1 year ago

18bc3543

Don't add `mhlo.layout_mode = "default"` since that is the default even in PJRT and will help reduce cruft in the IR

yashk2810 committed 1 year ago

a2973be0

[Pallas] Add runtime assert via checkify.check. This check will halt the TPU if triggered, meaning that we would need to restart the program to recover.

justinjfu committed 1 year ago

cff9e938

Deprecate a number of APIs in jax.lib.xla_client.

hawkinsp committed 1 year ago

e9c7ff0b

Remove deprecated API and migrate to new API.

ghpvnist committed 1 year ago

af50c212

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a684021c

Clone RandomAlgorithm into lax.py, instead of using the version from XLA.

hawkinsp committed 1 year ago

46f0a3ee

Merge pull request #24232 from ROCm:ci_rv_clang_clean

Google-ML-Automation committed 1 year ago

e4629f6a

#sdy add JAX Shardy support for memories.

bartchr808 committed 1 year ago

fb32841b

[pallas:mosaic_gpu] Added a test doing manual in kernel pipelining

superbobry committed 1 year ago

59ae2af6

[pallas:mosaic_gpu] Removed `mgpu.Barrier` sorting

superbobry committed 1 year ago

17248068

Merge pull request #24241 from hawkinsp:autodidax

Google-ML-Automation committed 1 year ago

bc3df0e3

Corresponding `build.py` updates after https://github.com/jax-ml/jax/commit/5132a188f75022a2ef1387d944f0285557876e24

ddunl committed 1 year ago

ee312afe

Merge pull request #24242 from ROCm:ci_bazel_build

Google-ML-Automation committed 1 year ago

e58ef1af

Port autodidax to use StableHLO instead of classic HLO.

hawkinsp committed 1 year ago

c0efa86b

[pallas:mosaic_gpu] `GPUBlockSpec` no longer accepts `swizzle`

superbobry committed 1 year ago

acd0e497

[Pallas TPU] Add lowering for scalar `jnp.log1p`

ayaka14732 committed 1 year ago

633cb315

[JAX] bazel build rocm changes

Ruturaj4 committed 1 year ago

89cd375c

[ROCm] Bring up clang support for JAX+XLA

Ruturaj4 committed 1 year ago

33bcd0cb

[Pallas] Add tests for scalar elementwise operations

ayaka14732 committed 1 year ago

3bd8ca48

Update Eigh kernel on GPU to use 64-bit interface when it is available.

dfm committed 1 year ago

6625a2b3

Merge pull request #24193 from dfm:vectorized-changelog

Google-ML-Automation committed 1 year ago

9d44d723

Fix listing of `vectorized` deprecation in changelog.

dfm committed 1 year ago

f55141ef

Refactor JAX's .bazelrc

nitins17 committed 1 year ago

5132a188

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

f8022e78

[sharding_in_types] Normalize partition specs when creating avals so that P(None, None) and P() are treated as replicated and equivalent. Shardings on avals are always normalized.

yashk2810 committed 1 year ago

8ef41a6e

Reenable some test cases that were disabled due to bugs that now seem fixed.

hawkinsp committed 1 year ago

66f52689

Move additional CI enabled/disabled configurations into jax BUILD files.

hawkinsp committed 1 year ago

19dbff53

Deprecate jax.lib.xla_client.PaddingType.

hawkinsp committed 1 year ago

aa3254d7

Older