Basic support of `jax.lax.gather` in TessellateIPU. (#23) #97

	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

Provide feedback