Add documentation to the repository

.github/workflows/doc-build.yml

@@ -0,0 +1,57 @@
+name: Build documentation
+
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - 'v[0-9]+.[0-9]+.[0-9]+'
+
+    paths:
+      - 'docs/source/**'
+      - 'docs/assets/**'
+      - 'optimum/**'
+      - '.github/workflows/doc-build.yml'
+  workflow_dispatch:
+
+jobs:
+  build_documentation:
+    runs-on: ubuntu-latest
+    env:
+      COMMIT_SHA: ${{ github.event.pull_request.head.sha }}
+      PR_NUMBER: ${{ github.event.number }}
+      EVENT_CONTEXT: ${{ toJSON(github.event) }}
+      PR_CLONE_URL: ${{ github.event.pull_request.head.repo.clone_url }}
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v4
+        with:
+          node-version: '18'
+          cache-dependency-path: "kit/package-lock.json"
+
+      - name: Set environment variables
+        run: |
+          cd optimum
+          version=`echo "$(grep '^__version__ =' tpu/version.py | cut -d '=' -f 2- | xargs)"`
+
+          if [[ $version == *.dev0 ]]
+          then
+            echo "VERSION=main" >> $GITHUB_ENV
+          else
+            echo "VERSION=v$version" >> $GITHUB_ENV
+          fi
+
+          cd ..
+
+      - name: Setup environment
+        run: |
+          pip install ".[quality]"
+
+      - name: Make documentation
+        shell: bash
+        run: |
+          doc-builder build optimum.tpu docs/source/ --repo_name optimum-tpu --build_dir tpu-doc-build/ --version ${{ env.VERSION }} --version_tag_suffix "" --html --clean
+          cd  tpu-doc-build/
+          mv optimum.tpu optimum-tpu
+          doc-builder push optimum-tpu --doc_build_repo_id "hf-doc-build/doc-build" --token "${{ secrets.HF_DOC_BUILD_PUSH }}" --commit_msg "Updated with commit $COMMIT_SHA See: https://github.com/huggingface/optimum-tpu/commit/$COMMIT_SHA" --n_retries 5

docs/source/howto/deploy.mdx

@@ -0,0 +1,82 @@
+# Deploying a Google TPU instance on Google Cloud Platform (GCP)
+
+
+## Context
+
+We assume the reader has already created a Google Cloud Platform (GCP) user or organisation account and an
+associated project.
+
+We also assume the reader to have the Google Cloud CLI installed. If not please follow the links right after to
+[install](https://cloud.google.com/sdk/docs/install) and [setup](https://cloud.google.com/sdk/docs/initializing).
+
+## Creating the initial TPUVM on GCP
+
+In order to create your initial TPU instance, you will need to provide some information:
+
+- The zone in GCP you would like to see the instance being deployed (close to the reader for development purpose, close to the end user for production for instance)
+- Which kind of TPU you would like to target
+- Which version of the TPU runtime you want to leverage on the instance
+- Custom instance name to quickly skim and refer back to the instance
+
+Overall the end command looks like this:
+
+```bash
+gcloud compute tpus tpu-vm create <ref_instance_name> \
+--zone=<deploiment_zone> \
+--accelerator-type=<target_tpu_generation> \
+--version=<runtime_version>
+```
+
+### Deploying a TPU v5litepod-8 instance
+
+In our case we will be deploying a `v5litepod-8` instance name `optimum-tpu-get-started`
+in the GCP region `us-west4-a` using the latest `v2-alpha-tpuv5-lite` runtime version.
+
+Of course, feel free to adjust all these parameters to the one that match with your usage and quotas.
+
+Before creating the instance, please make sure to install `gcloud alpha component` as it is required to be able to
+target TPUv5 VMs: `gcloud components install alpha`
+
+```bash
+gcloud alpha compute tpus tpu-vm create optimum-tpu-get-started \
+--zone=us-west4-a \
+--accelerator-type=v5litepod-8 \
+--version=v2-alpha-tpuv5
+```
+
+## Connecting to the instance
+
+```bash
+gcloud compute tpus tpu-vm ssh <ref_instance_name> --zone=<deploiment_zone>
+$ >
+```
+
+In the example above deploying v5litepod-8 it would be something like:
+
+```bash
+gcloud compute tpus tpu-vm ssh optimum-tpu-get-started --zone=us-west4-a
+$ >
+```
+
+## Setting up the instance to run AI workloads on TPUs
+
+### Optimum-TPU with PyTorch/XLA
+
+If you want to leverage PyTorch/XLA through Optimum-TPU, it should be as simple as
+
+```bash
+$ python3 -m pip install optimum-tpu
+$ export PJRT_DEVICE=TPU
+```
+
+Now you can validate the installation with the following command which should print `xla:0` as we do have a single
+TPU device bound to this instance.
+
+```bash
+$ python -c "import torch_xla.core.xla_model as xm; print(xm.xla_device())"
+xla:0
+```
+
+### Optimum-TPU with JAX
+
+JAX is coming very soon - stay tuned!

docs/source/howto/overview.mdx

@@ -0,0 +1,25 @@
+# Optimum-TPU How To
+
+This page gives you access to handful of walkthrough scenarios to leverage Google TPUs for your use case.
+
+## Looking for?
+
+<div class="mt-10">
+    <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-2 md:gap-y-4 md:gap-x-5">
+        <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./serving.mdx">
+            <div class="w-full text-center bg-gradient-to-r from-slate-900 to-slate-700 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">
+                Deploying Google Cloud TPU instance
+            </div>
+            <p class="text-gray-700">
+
+            </p>
+        </a>
+        <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./deploy.mdx">
+            <div class="w-full text-center bg-gradient-to-r from-slate-900 to-slate-700 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">
+                Deploying a Text-Generation-Inference server on Google Cloud TPU instance
+            </div>
+            <p class="text-gray-700">
+
+            </p>
+        </a>
+</div>

docs/source/index.mdx

@@ -0,0 +1,69 @@
+<!---
+Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+# 🤗 Optimum TPU
+
+Optimum TPU provides all the necessary machinery to leverage and optimize AI workloads runningo on [Google Cloud TPU devices](https://cloud.google.com/tpu/docs).
+
+The API provides the overall same user-experience as Hugging Face transformers with the minimum amount of changes required to target performance for both inference.
+
+Training support is underway, stay tuned! 🚀
+
+
+## Installation
+
+Optimum TPU is meant to reduce as much as possible the friction in order to leverage Google Cloud TPU accelerators.
+As such, we provide a pip installable package to make sure everyone can get easily started.
+
+### Run Cloud TPU with pip
+```bash
+pip install optimum-tpu
+```
+
+### Run Cloud TPU within Docker container
+
+### PyTorch
+```bash
+export TPUVM_IMAGE_URL=us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla
+export TPUVM_IMAGE_VERSION=8f1dcd5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526
+docker pull
+docker run -ti --rm --privileged --network=host ${TPUVM_IMAGE_URL}@sha256:${TPUVM_IMAGE_VERSION} bash
+```
+
+From there you can install optimum-tpu through the pip instructions above.
+
+
+<div class="mt-10">
+    <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-2 md:gap-y-4 md:gap-x-5">
+        <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./howto/overview">
+            <div class="w-full text-center bg-gradient-to-br from-red-500 to-red-800 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">
+                How-to guides
+            </div>
+            <p class="text-gray-700">
+
+            </p>
+        </a>
+        <a
+            class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg"
+            href="./package_reference/trainer"
+        >
+            <div class="w-full text-center bg-gradient-to-br from-green-500 to-green-800 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">
+                Reference
+            </div>
+            <p class="text-gray-700">Technical descriptions of how the classes and methods of Optimum TPU</p>
+        </a>
+    </div>
+</div>