Add new CI scripts as part of JAX CI Rework #59

	name: Run Pytest TPU tests

	on:
	pull_request:
	branches:
	- main

	jobs:
	build:
	strategy:
	matrix:
	runner: ["linux-x86-ct5lp-224-8tpu"]
	python: ["3.10"]

	runs-on: ${{ matrix.runner }}
	container:
	image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"

	env:
	# GitHub actions run in Docker by defaut. Disable running the `setup_docker.sh` script.
	JAXCI_RUN_DOCKER_CONTAINER: 0
	JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}

	steps:
	- uses: actions/checkout@v3
	# Halt for testing
	- name: Wait For Connection
	uses: ./actions/ci_connection/
	- name: Build jaxlib
	run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib.env"
	- name: Install pytest
	env:
	JAXCI_PYTHON: python${{ matrix.python }}
	run: $JAXCI_PYTHON -m pip install pytest
	- name: Install Test requirements
	env:
	JAXCI_PYTHON: python${{ matrix.python }}
	run: \|
	$JAXCI_PYTHON -m pip install -r build/test-requirements.txt
	$JAXCI_PYTHON -m pip install -r build/collect-profile-requirements.txt
	- name: Install Libtpu
	env:
	JAXCI_PYTHON: python${{ matrix.python }}
	run: $JAXCI_PYTHON -m pip install --pre libtpu-nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
	- name: Run Pytest TPU tests
	run: ./ci/run_pytest.sh "ci/envs/run_tests/pytest_tpu.env"

Provide feedback