End-to-end tests for torch_xla llama and mixtral #197

	name: CPU tests

	on:
	push:
	branches:
	- main
	pull_request:
	schedule: # Schedule the job run at 12AM PST daily.
	- cron: "0 8 * * *"

	jobs:
	pytest:
	runs-on: ubuntu-22.04
	strategy:
	matrix:
	python-version: ["3.10", "3.11"]
	container:
	image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_${{ matrix.python-version }}_tpuvm
	steps:
	- uses: actions/checkout@v4
	- name: Install torchax
	run: \|
	pip install 'torchax @ git+https://[email protected]/pytorch/xla.git#subdirectory=torchax'
	- name: Install dev dependencies
	run: \|
	python -m pip install --upgrade pip
	pip install -e '.[dev]'
	- name: Run PyTest
	env:
	# TODO(https://github.com/AI-Hypercomputer/torchprime/issues/14): Remove and burn the token.
	HF_TOKEN: ${{ secrets.HF_TOKEN }}
	run: \|
	export PJRT_DEVICE=CPU
	export JAX_PLATFORMS=cpu
	pytest -v -m "not deepseek"

	# TODO(https://github.com/AI-Hypercomputer/torchprime/issues/75): Fix the failure on
	# torch 2.6, then remove the requirements and extra install.
	pip install --force -r torchprime/experimental/torchax_models/deepseek_v3/requirements.txt
	pytest -v -m deepseek torchprime/experimental/torchax_models/deepseek_v3/tests

Provide feedback