Replace deprecated `jax.tree_` functions with `jax.tree.` #915

	name: CI

	on:
	# Triggers the workflow on push or pull request events but only for the main branch
	push:
	branches: [ main ]
	pull_request:
	branches: [ main ]
	# Allows you to run this workflow manually from the Actions tab
	workflow_dispatch:

	jobs:
	build:
	runs-on: ubuntu-latest
	strategy:
	fail-fast: false
	matrix:
	python-version: [3.7, 3.9]
	steps:
	- name: Cancel previous
	uses: styfle/[email protected]
	with:
	access_token: ${{ github.token }}
	if: ${{github.ref != 'refs/head/main'}}
	- uses: actions/checkout@v2
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@v2
	with:
	python-version: ${{ matrix.python-version }}
	- name: Get pip cache dir
	id: pip-cache
	run: \|
	python -m pip install --upgrade pip wheel
	echo "::set-output name=dir::$(pip cache dir)"
	- name: pip cache
	uses: actions/cache@v2
	with:
	path: ${{ steps.pip-cache.outputs.dir }}
	key: ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}
	- name: Install dependencies
	run: \|
	pip install jax jaxlib
	- name: Test with pytest
	# validation_test is too slow (it needs to run on TPUs)
	# time_stepping_test requires reducing tolerance too much in GitHub actions (maybe a JAX bug?)
	run: \|
	pip install -e ".[complete]"
	pytest -n auto --forked jax_cfd --ignore=jax_cfd/base/validation_test.py --ignore=jax_cfd/spectral/time_stepping_test.py

Provide feedback