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

tweak readme pmap imports ()

Browse Source
This commit is contained in:
brett koonce 2020-02-20 18:44:21 -06:00 committed by GitHub
parent 218a1711d2
commit 8372a70079
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions

3
README.md

@ -272,7 +272,8 @@ replicated and executed in parallel accross devices.
Here's an example on an 8-GPU machine:
```python
from jax import random
from jax import random, pmap
import jax.numpy as np
# Create 8 random 5000 x 6000 matrices, one per GPU
keys = random.split(random.PRNGKey(0), 8)