ROCm/jax

Pull Requests Commits

Remove _manual_axes from NamedSharding since we can now track the manual axes on the mesh.

yashk2810 committed 1 year ago

0a9d0bec

Merge pull request #28077 from hawkinsp:postrelease2

Google-ML-Automation committed 1 year ago

b0b395f5

Merge pull request #28078 from justinjfu:pipe_docs_v2

Google-ML-Automation committed 1 year ago

8b2f04fc

[Pallas] Fix relative links in pipelining docs

justinjfu committed 1 year ago

cbcf3ab6

Update version numbers after release.

hawkinsp committed 1 year ago

fbd6db3f

Merge branch 'release/0.6.0'

hawkinsp committed 1 year ago

407e0ef4

Merge pull request #28069 from dfm:fix-argums-partial

Google-ML-Automation committed 1 year ago

003713cc

Return False in `is_env_present` if importing kubernetes leads to a ModuleNotFoundError

yashk2810 committed 1 year ago

127aa762

Return False in `is_env_present` if importing kubernetes leads to a ModuleNotFoundError

yashk2810 committed 1 year ago

a31e53a6

Fix a bug in reduce_window sharding rule where padding is a tuple but we were checking for a scalar instead. Fixes https://github.com/jax-ml/jax/issues/28070

yashk2810 committed 1 year ago

fb16d533

JAX release v0.6.0

hawkinsp committed 1 year ago

b3d5085a

Fix a bug in reduce_window sharding rule where padding is a tuple but we were checking for a scalar instead. Fixes https://github.com/jax-ml/jax/issues/28070

yashk2810 committed 1 year ago

5f6b99a1

Merge pull request #28072 from justinjfu:pipe_docs_v2

Google-ML-Automation committed 1 year ago

af37883c

[Pallas] Move pipelining doc figures to static folder

justinjfu committed 1 year ago

b61f87f5

Fix bug in argnums_partial_except when static_argnums is unsorted.

dfm committed 1 year ago

9afc047b

Merge pull request #28062 from jax-ml:sharded-computation

Google-ML-Automation committed 1 year ago

b6bcfb63

Add explicit sharding to Introduction to parallel programming doc

yashk2810 committed 1 year ago

4ad5dd83

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7f7b5c6e

Merge pull request #28052 from justinjfu:pipe_docs_v2

Google-ML-Automation committed 1 year ago

310fafa9

Update job name to show "build only" on Linux ARM64.

nitins17 committed 1 year ago

8073d78a

[Pallas] Refactor pipelining docs.

justinjfu committed 1 year ago

cd5d4821

Merge pull request #28018 from Cjkkkk:disable_packed_layout_at_ampere

Google-ML-Automation committed 1 year ago

74f1d887

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7d66cdd3

Merge pull request #27994 from vfdev-5:fix-313-suppression

Google-ML-Automation committed 1 year ago

bd88eeb5

Merge pull request #28004 from apaszke:mgpu-layouts-docs

Google-ML-Automation committed 1 year ago

a9f6d3a9

Add basic documentation for Pallas:MGPU layouts and transforms

apaszke committed 1 year ago

dabe7b7e

[pallas] Use `*Params` dataclasses to pass `compiler_params=` to `pl.pallas_call`

superbobry committed 1 year ago

19e2315b

Merge pull request #28056 from apaszke:mgpu-memspace-picture

Google-ML-Automation committed 1 year ago

f8e26bdb

[Pallas:MGPU] Add a diagram outlining GPU memory spaces

apaszke committed 1 year ago

98a9ef44

Merge pull request #27983 from apaszke:mgpu-memory-space-docs

Google-ML-Automation committed 1 year ago

02bdffd0

Newer Older