Fix import bug issue with `jax.stages.ArgInfo` in past JAX versions. … #374

	name: 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:
	lint_and_typecheck:
	runs-on: ubuntu-latest
	timeout-minutes: 10
	steps:
	- name: Cancel previous
	uses: styfle/cancel-workflow-action@0.11.0
	with:
	access_token: ${{ github.token }}
	if: ${{github.ref != 'refs/head/main'}}
	- uses: actions/checkout@v3
	- name: Set up Python 3.10
	uses: actions/setup-python@v4
	with:
	python-version: "3.10"
	- uses: pre-commit/action@v3.0.0

	unit_tests:
	runs-on: ubuntu-latest
	timeout-minutes: 10
	steps:
	- name: Cancel previous
	uses: styfle/cancel-workflow-action@0.11.0
	with:
	access_token: ${{ github.token }}
	if: ${{github.ref != 'refs/head/main'}}
	- uses: actions/checkout@v3
	- name: Update pip
	id: pip-cache
	run: \|
	python3 -m pip install --upgrade pip
	- name: Local install & test requirements
	run: \|
	pip3 install -e ./
	pip3 install -r ./test-requirements.txt
	# Run repository unit tests on latest JAX
	- name: Run unit tests JAX latest
	run: \|
	pytest --tb=short -v --log-cli-level=INFO ./
	- name: JAX 0.3.16 installation
	run: \|
	pip3 install numpy==1.24.3 scipy==1.10.1
	pip3 install chex==0.1.6 jax==0.3.16 jaxlib==0.3.15 -f https://storage.googleapis.com/jax-releases/jax_releases.html
	# Run repository unit tests on JAX 0.3
	- name: Run unit tests JAX 0.3.16
	run: \|
	pytest --tb=short -v --log-cli-level=INFO ./

Provide feedback