ROCm/jax

Pull Requests Commits

[ROCm] Fix README.md to update AMD JAX installation instructions

Ruturaj4 committed 2 years ago

3c3b08df

Merge pull request #24290 from jax-ml:dependabot/github_actions/actions/upload-artifact-4.4.3

Google-ML-Automation committed 2 years ago

90cd8a79

Merge pull request #24289 from jax-ml:dependabot/github_actions/actions/cache-4.1.1

Google-ML-Automation committed 2 years ago

13bc4978

Bump actions/upload-artifact from 4.4.1 to 4.4.3

dependabot[bot] committed 2 years ago

Verified 93adc0e9

Bump actions/cache from 4.1.0 to 4.1.1

dependabot[bot] committed 2 years ago

Verified 0fdd6535

[Shardy] Inline meshes when using shardy and get rid of global meshes from the MLIR body.

yashk2810 committed 2 years ago

824ccd71

#sdy Run inlined mesh lifter pass at the end of JAX lowering.

bartchr808 committed 2 years ago

75e22f2c

Merge pull request #24271 from jakevdp:hist-doc

Google-ML-Automation committed 2 years ago

d15d70d6

Merge pull request #24264 from ROCm:ci_apt_update

Google-ML-Automation committed 2 years ago

1de9f25c

Merge pull request #24274 from ROCm:ci_linalg_fix

Google-ML-Automation committed 2 years ago

57ef7a4a

Port Schur Decomposition to XLA's FFI

Paweł Paruzel committed 2 years ago

23fdb912

Improve documentation for jnp.histogram & related APIs

jakevdp committed 2 years ago

9b9fc5af

Port Tridiagonal Reduction to XLA's FFI

Paweł Paruzel committed 2 years ago

ec68d420

[ROCm] jaxlib linalg fix

Ruturaj4 committed 2 years ago

ee223d40

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

b6f38bcc

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

44cfbbe3

[sharding_in_types] Add constraints during lowering for dot_general and reduce_sum so that we can enforce the sharding we choose during tracing

yashk2810 committed 2 years ago

4be1e332

Fix repr of sharding in aval when a dimension is sharded on multiple mesh axes

yashk2810 committed 2 years ago

8139c531

[sharding_in_types] Add sharding rule for reduce sum which is just drop the specs for the axis we are reducing over

yashk2810 committed 2 years ago

5b8775dc

Better repr of aval when shardings are present

yashk2810 committed 2 years ago

89fcd9f1

[ROCm] apt update

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

cff9e938

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

hawkinsp committed 2 years ago

e9c7ff0b

Remove deprecated API and migrate to new API.

ghpvnist committed 2 years ago

af50c212

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

a684021c

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

hawkinsp committed 2 years ago

46f0a3ee

Merge pull request #24232 from ROCm:ci_rv_clang_clean

Google-ML-Automation committed 2 years ago

e4629f6a

#sdy add JAX Shardy support for memories.

bartchr808 committed 2 years ago

fb32841b

Older