Fix pytype and clean up types across codebase. #1215

	name: Unit tests

	on: [push, pull_request]

	jobs:
	unittests:
	name: "unittests (${{ matrix.python-version }})"
	runs-on: ubuntu-latest
	strategy:
	matrix:
	python-version: ["3.9", "3.10", "3.11"]

	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
	run: \|
	python -m pip install --upgrade pip
	pip install -r requirements/requirements.txt \
	-r requirements/requirements-test.txt \
	-r requirements/requirements-examples.txt

	- name: Run pytype
	run: \|
	pytype rax/ examples/flax_integration/ examples/approx_metrics/

	- name: Run tests
	run: \|
	python -m unittest discover -v -s ./rax -p "*_test.py"
	python -m unittest discover -v -s ./examples/flax_integration -p "*_test.py"
	python -m unittest discover -v -s ./examples/approx_metrics -p "*_test.py"

Provide feedback