update installation docs

brainpy · Mar 5, 2024 · 0e866f4 · 0e866f4
1 parent b67dcea
commit 0e866f4
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/docs/index.rst b/docs/index.rst
@@ -19,16 +19,24 @@ Installation
 
           pip install -U brainpy[cpu]  # windows, linux, macos
 
-    .. tab-item:: GPU (CUDA)
+    .. tab-item:: GPU (CUDA 11.0)
 
        .. code-block:: bash
 
-          # for CUDA 11.0, Linux only
           pip install -U brainpy[cuda11] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
-          # for CUDA 12.0, Linux only
+    .. tab-item:: GPU (CUDA 12.0)
+
+       .. code-block:: bash
+
           pip install -U brainpy[cuda12] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
+    .. tab-item:: TPU
+
+       .. code-block:: bash
+
+          pip install -U brainpy[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+
 
 For more information, please see `installation <quickstart/installation.html>`_ section.
 

diff --git a/docs/quickstart/installation.rst b/docs/quickstart/installation.rst
@@ -17,8 +17,7 @@ To install brainpy with minimum requirements (has installed ``jax`` and ``jaxlib
 
 .. code-block:: bash
 
-    pip install brainpy # for CPU
-
+    pip install brainpy
 
 
 Minimum requirements (with dependencies)
@@ -35,6 +34,9 @@ To install brainpy with minimum requirements (only depends on ``jax``), you can
     pip install brainpy[cuda11_mini] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  # for CUDA 11.0
     pip install brainpy[cuda12_mini] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  # for CUDA 12.0
 
+    # or
+
+    pip install brainpy[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html  # for google TPU
 
 
 CPU with all dependencies
@@ -61,6 +63,18 @@ To install a GPU-only version of BrainPy, you can run
 
 
 
+TPU with all dependencies
+-------------------------
+
+BrainPy supports Google Cloud TPU. To install BrainPy along with appropriate versions of jax,
+you can run the following in your cloud TPU VM:
+
+.. code-block:: bash
+
+    pip install brainpy[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html  # for google TPU
+
+
+
 ``brainpylib``
 --------------