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

Merge pull request #13483 from google:cross

Browse Source 
PiperOrigin-RevId: 492314228
This commit is contained in:
jax authors 2022-12-01 15:23:19 -08:00
commit a6c7fa85bb
3 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docs/jax_array_migration.md
View File

@ -1,3 +1,4 @@
(jax-array-migration)=
# jax.Array migration
**yashkatariya@**
@ -17,6 +18,8 @@ the unified jax.Array
After the migration is complete `jax.Array` will be the only type of array in
JAX.
This doc explains how to migrate existing codebases to `jax.Array`. For more information on using `jax.Array` and JAX parallelism APIs, see the [Parallelism with JAX](https://jax.readthedocs.io/en/latest/notebooks/Parallelism_with_JAX.html) tutorial.
### How to enable jax.Array?

2
docs/notebooks/Parallelism_with_JAX.ipynb
View File

@ -17,6 +17,8 @@
"source": [
"**This tutorial discusses parallelism via `jax.Array`, the unified array object model available in JAX v0.4.0 and newer.**\n",
"\n",
"See {ref}`jax-array-migration` guide for migrating existing pre-v0.4.0 codebases to `jax.Array`.\n",
"\n",
"**The features required by `jax.Array` are not supported by the Colab TPU runtime at this time.**"
]
},

2
docs/notebooks/Parallelism_with_JAX.md
View File

@ -19,6 +19,8 @@ kernelspec:
**This tutorial discusses parallelism via `jax.Array`, the unified array object model available in JAX v0.4.0 and newer.**
See {ref}`jax-array-migration` guide for migrating existing pre-v0.4.0 codebases to `jax.Array`.
**The features required by `jax.Array` are not supported by the Colab TPU runtime at this time.**
```{code-cell}