Basic support of jax.lax.gather in TessellateIPU.
#94
Workflow file for this run
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: notebook-tests | |
| env: | |
| GIT_MAIN_BRANCH: "main" | |
| # Controls when the workflow will run. | |
| on: | |
| push: | |
| branches: [ "main" ] | |
| pull_request: | |
| branches: [ "main" ] | |
| # Allows you to run this workflow manually from the Actions tab. | |
| workflow_dispatch: | |
| jobs: | |
| # Tessellate unit tests using IPU model | |
| tessellate_notebook_tests_ipu_model: | |
| if: github.repository == 'graphcore-research/tessellate-ipu' | |
| runs-on: ubuntu-latest | |
| container: graphcore/pytorch:3.2.0-ubuntu-20.04 | |
| timeout-minutes: 20 | |
| steps: | |
| - name: Cancel previous | |
| uses: styfle/[email protected] | |
| with: | |
| access_token: ${{ github.token }} | |
| if: ${{github.ref != 'refs/head/main'}} | |
| - uses: actions/checkout@v3 | |
| - name: Update pip & install Jupyter | |
| id: pip-cache | |
| run: | | |
| python3 -m pip install --upgrade pip wheel setuptools jupyter | |
| python3 -m pip install --upgrade build cmake nanobind ninja pip scikit_build_core[pyproject] | |
| # Install JAX for IPU. | |
| - name: Install JAX for IPU | |
| run: | | |
| pip3 install --default-timeout=100 -U numpy==1.23.5 scipy etils pytest | |
| pip3 install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk320 -f https://graphcore-research.github.io/jax-experimental/wheels.html | |
| # Build & install sdist package. | |
| - name: Install TessellateIPU | |
| run: | | |
| python3 -m build --no-isolation | |
| pip3 install dist/*.tar.gz | |
| # Run TessellateIPU notebooks. | |
| - name: Run TessellateIPU notebooks (IPU model) | |
| run: | | |
| python3 .github/scripts/run_ipu_notebook.py ./notebooks/01-tessellate-ipu-tile-api-basics.ipynb |