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

	name: Build and minimal test

	on:
	push:
	branches:
	- main
	pull_request:
	branches:
	- main

	jobs:
	build:

	runs-on: ubuntu-latest
	strategy:
	matrix:
	python-version: ['3.8', '3.9', '3.10']

	steps:
	- 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
	# Use `install -e` for local fedjax to prevent actual fedjax files from
	# being packaged into cache.
	run: \|
	python -m pip install --upgrade pip setuptools wheel
	pip install -e .
	- name: Basic test
	# We run tests using `python -I` to prevent the script directory to be
	# added to sys.path.
	run: \|
	python -I fedjax/fedjax_test.py
	- name: Main package test
	run: \|
	./run_tests.sh
	- name: Examples test
	run: \|
	pushd examples
	python fed_avg_test.py
	popd
	- name: Build documentation
	run: \|
	pip install -r docs/requirements.txt
	sphinx-build -M html docs docs/_build

Provide feedback