Fix `device` for JAX #236

	name: CI

	on:
	- push
	- pull_request

	jobs:
	test:
	runs-on: ubuntu-latest
	strategy:
	matrix:
	python-version: [3.8, 3.9, "3.10", "3.11"]
	steps:
	- uses: actions/checkout@v2
	with:
	fetch-depth: 0
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@v2
	with:
	python-version: ${{ matrix.python-version }}
	- name: Install dependencies
	run: \|
	sudo apt-get install gfortran
	pip install --upgrade pip setuptools 'setuptools_scm[toml]' setuptools_scm_git_archive numpy Cython
	python setup.py --version
	LAB_BUILD=1 pip install --no-cache-dir -U -r requirements.txt \| cat
	pip install --upgrade numpy
	- name: Test
	run: \|
	JAX_ENABLE_X64=1 pytest -v --cov=lab --cov-report term-missing
	coveralls --service=github
	env:
	GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
	COVERALLS_FLAG_NAME: ${{ matrix.test-name }}
	COVERALLS_PARALLEL: true

	finish:
	name: Finish Coveralls
	needs: test
	runs-on: ubuntu-latest
	steps:
	- name: Finish Coveralls
	uses: coverallsapp/github-action@v1
	with:
	github-token: ${{ secrets.github_token }}
	parallel-finished: true

Provide feedback