Additional corrections to IPU documentation. (#11)

Author: balancap

README.md

@@ -84,7 +84,11 @@ Alternatively, like other JAX flags, these can be set using environment variable
 
 ## License
 
-The project remains licensed under the **Apache License 2.0**. No additional Python or C++ dependency has been introduced compared to the original JAX source code (outside Graphcore Poplar and Poplibs licensed under MIT license).
-* JAX license: https://github.com/graphcore-research/jax-experimental/LICENSE
-* JAXLIB license (including compiled dependencies): https://github.com/graphcore-research/jax-experimental/build/LICENSE.txt
-* Graphcore license agreements: https://docs.graphcore.ai/en/latest/licenses.html
+The project remains licensed under the **Apache License 2.0**, with the following files unchanged:
+* [JAX Apache license 2.0](./LICENSE);
+* [JAXLIB and compiled dependencies licenses;](build/LICENSE.txt)
+
+The additional dependencies introduced for Graphcore IPU support are:
+* [Poplar: licensed under the Graphcore EULA](https://docs.graphcore.ai/projects/eula/en/latest/);
+* [PopLibs: licensed under MIT license](https://docs.graphcore.ai/en/latest/licenses.html);
+* [TensorFlow XLA for IPU: licensed under Apache license 2.0](https://docs.graphcore.ai/en/latest/licenses.html);

WORKSPACE

@@ -21,8 +21,8 @@ load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 git_repository(
     name = "org_tensorflow",
     # Jaxlib 0.3.15 compatible with JAX 0.3.16
-    branch = "jaxlibv0.3.15/sdk-release-3.0-lean-jax",
-    remote = "git@github.com:graphcore-research/tensorflow-jax-mk2-experimental.git"
+    branch = "jaxlibv0.3.15/sdk-release-3.1-lean-jax",
+    remote = "git@github.com:graphcore-research/tensorflow-jax-experimental.git"
 )
 
 # For development, one often wants to make changes to the TF repository as well

ipu/docs/build.md

@@ -1,21 +1,21 @@
-# How to build experimental JAX on IPU Mk2
+# How to build experimental JAX on IPU
 
-Building and testing JAX on IPU requires the following base configuration: Ubuntu 20.04, [Graphcore Poplar SDK 3.1](https://www.graphcore.ai/posts/poplar-sdk-3.1-now-available) and [Bazel 5.2.0](https://docs.bazel.build/versions/5.2.0/install.html). Other OS or Python versions have not been tested.
+Building and testing JAX on IPU requires the following base configuration: Ubuntu 20.04, Python 3.8, [Graphcore Poplar SDK 3.1](https://www.graphcore.ai/posts/poplar-sdk-3.1-now-available) and [Bazel 5.2.0](https://docs.bazel.build/versions/5.2.0/install.html). Other OS or Python versions have not been tested.
 
-Additionally, one needs to install to the following packages in your Python environment:
+Additionally, you need to install to the following packages in your Python environment:
 ```bash
 pip install numpy==1.21.6 scipy cython pytest
 ```
-Note: the build process will work with more recent versions of Numpy, but that will then limit Numpy backward compatibility.
+Note: the build process will work with more recent versions of NumPy, but that will then limit NumPy backward compatibility.
 
 ## Build IPU JAXLIB
 
 Building `jaxlib` currently requires the branch `jax-v0.3.15-ipu`. Once the branch checked out, the build process is similar to other backends:
 ```bash
 export TF_POPLAR_BASE=#...poplar install directory
-python build/build.py --bazel_options=--override_repository=org_tensorflow=PATH/tensorflow-jax-mk2-experimental
+python build/build.py --bazel_options=--override_repository=org_tensorflow=PATH/tensorflow-jax-experimental
 ```
-The `override_repository` config is optional. By default, the build process will pull the experimental IPU TensorFlow XLA code from the repository https://github.com/graphcore-research/tensorflow-jax-mk2-experimental.
+The `override_repository` config is optional. By default, the build process will pull the experimental IPU TensorFlow XLA code from the repository https://github.com/graphcore-research/tensorflow-jax-experimental.
 
 If the build is successful, a binary `jaxlib` Python wheel will be produced in the `dist/` directory.