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

Update JAX official doc: point out that the device numbers are not in numerical order because of the underlying torus hardware topology.

Browse Source 
PiperOrigin-RevId: 625401373
This commit is contained in:
Yue Sheng 2024-04-16 11:36:18 -07:00 committed by jax authors
parent e83b0ce3f2
commit 1a650cdc00
2 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
View File

@ -412,6 +412,8 @@
"id": "uRLpOcmNj_Vt"
},
"source": [
"The device numbers here are not in numerical order, because the mesh reflects the underlying toroidal topology of the device.\n",
"\n",
"By writing `PositionalSharding(ndarray_of_devices)`, we fix the device order and the initial shape. Then we can reshape it:"
]
},

2
docs/notebooks/Distributed_arrays_and_automatic_parallelization.md
View File

@ -192,6 +192,8 @@ sharding
+++ {"id": "uRLpOcmNj_Vt"}
The device numbers here are not in numerical order, because the mesh reflects the underlying toroidal topology of the device.
By writing `PositionalSharding(ndarray_of_devices)`, we fix the device order and the initial shape. Then we can reshape it:
```{code-cell}