mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #22887 from justinjfu:pallas_distr_docs

Browse Source 
PiperOrigin-RevId: 661368285
This commit is contained in:
jax authors 2024-08-09 12:47:48 -07:00
commit 35c2454ea6
13 changed files with 3282 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/_static/pallas/distributed/all_gather.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 58 KiB

1
docs/_static/pallas/distributed/race_condition.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 101 KiB

1
docs/_static/pallas/distributed/rdma_recv.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 63 KiB

1
docs/_static/pallas/distributed/rdma_send.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 64 KiB

1
docs/_static/pallas/distributed/rdma_start.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 63 KiB

1
docs/_static/pallas/distributed/reduce_scatter_1.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 46 KiB

1
docs/_static/pallas/distributed/reduce_scatter_2.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 64 KiB

1
docs/_static/pallas/distributed/reduce_sum_1.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 18 KiB

1
docs/_static/pallas/distributed/reduce_sum_2.svg vendored Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 94 KiB

2
docs/conf.py
View File

@ -132,6 +132,7 @@ exclude_patterns = [
'notebooks/*.md',
'pallas/quickstart.md',
'pallas/tpu/pipelining.md',
'pallas/tpu/distributed.md',
'pallas/tpu/matmul.md',
'jep/9407-type-promotion.md',
'autodidax.md',
@ -221,6 +222,7 @@ nb_execution_excludepatterns = [
# Requires accelerators
'pallas/quickstart.*',
'pallas/tpu/pipelining.*',
'pallas/tpu/distributed.*',
'pallas/tpu/matmul.*',
'sharded-computation.*',
'distributed_data_loading.*'

1743
docs/pallas/tpu/distributed.ipynb Normal file
View File

File diff suppressed because it is too large Load Diff

1527
docs/pallas/tpu/distributed.md Normal file
View File

File diff suppressed because it is too large Load Diff

1
docs/pallas/tpu/index.rst
View File

@ -9,3 +9,4 @@ TPU specific documentation.
details
pipelining
matmul
distributed