Crash fix for latest JAX + new release. #154

	# Copyright 2023 Google LLC
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	name: Run tests

	on:
	pull_request:

	jobs:
	run-test:
	strategy:
	matrix:
	python-version: [ 3.9, 3.11 ]
	os: [ ubuntu-latest ]
	fail-fast: false
	runs-on: ${{ matrix.os }}
	steps:
	- name: Checkout code
	uses: actions/checkout@v2

	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@v2
	with:
	python-version: ${{ matrix.python-version }}

	- name: Install dependencies
	run: \|
	python -m pip install --upgrade pip
	python -m pip install -r ./tests/requirements.txt

	- name: Checks with pre-commit
	uses: pre-commit/[email protected]

	- name: Test with pytest
	run: \|
	python -m pip install .
	python -m tests

Provide feedback