build_tensorflow_tpu.sh

#!/bin/bash
set -e
set -v
sudo rm -rf ~/.cache/bazel  # clear bazel cache (see tensorflow#55563)
export TENSORFLOW_SHA=0d5668cbdc6b46d099bd3abd93374c09b2e8121f
wget https://github.com/tensorflow/tensorflow/archive/$TENSORFLOW_SHA.tar.gz
tar zxf $TENSORFLOW_SHA.tar.gz
rm -f $TENSORFLOW_SHA.tar.gz
mv tensorflow-$TENSORFLOW_SHA tensorflow
cd tensorflow
export BAZEL_VERSION=`cat .bazelversion`
sudo mkdir -p /usr/local/lib/bazel/bin
sudo wget -nc -P /usr/local/lib/bazel/bin https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-linux-x86_64
sudo chmod +x /usr/local/lib/bazel/bin/bazel-$BAZEL_VERSION-linux-x86_64
rm -rf ~/.venv310
python3.10 -m venv ~/.venv310
. ~/.venv310/bin/activate
pip install -U wheel packaging
pip install -U keras_preprocessing --no-deps
pip install "jax[tpu]==0.3.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
yes '' | python configure.py  # use the default settings
bazel build --config=tpu //tensorflow/tools/pip_package:build_pip_package
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
ls /tmp/tensorflow_pkg/tensorflow-*.whl

Build_TF_TPU.md

Tensorflow 2.10.0 wheel, built for Python 3.10 with TPU support, compactible with JAX v0.3.6, jaxlib v0.3.5 and libtpu-nightly v0.1.dev20220412.

Get the required Tensorflow commit of jaxlib v0.3.5 from https://github.com/google/jax/blob/jaxlib-v0.3.5/WORKSPACE#L25.

Build: Run the build script.

tensorflow-2.10.0-cp310-cp310-linux_x86_64.whl.0

tensorflow-2.10.0-cp310-cp310-linux_x86_64.whl.1