Use pytest to run tests under CI.

We cannot use bazel until bazelbuild/bazel#7091 is resolved since bazel invokes
Python tests via `python path/to/test.py` and this causes python to add
`path/to` as the first entry in sys.path. This means for any test in
`haiku/_src` we have a collision between `typing` in the standard library and
`haiku/_src/typing.py` causing tests to fail.

Running Haiku tests by passing modules to python (e.g. `python -m path.to.test`)
works, as does running via `pytest`. We prefer `pytest` since it discovers tests
and has support for running tests in parallel.

PiperOrigin-RevId: 291146369
Change-Id: I01165f0d0dc20f1c520c02ae9b1da155dbbf0eb6

requirements-test.txt

@@ -2,3 +2,4 @@ cloudpickle>=1.2.2
 mock>=3.0.5
 tensorflow-datasets>1
 docutils
+pytest-xdist

test.sh

@@ -21,32 +21,24 @@ set -x
 
 virtualenv -p python3 .
 source bin/activate
-python3 --version
+python --version
 
-# Run setup.py, this installs the python dependencies
-python3 setup.py install
-
-N_JOBS=$(grep -c ^processor /proc/cpuinfo)
+# Install JAX.
+python -m pip install -r requirements-jax.txt
+python -c 'import jax; print(jax.__version__)'
 
-echo ""
-echo "Bazel will use ${N_JOBS} concurrent job(s)."
-echo ""
+# Run setup.py, this installs the python dependencies
+python -m pip install .
 
 # Python test dependencies.
-python3 -m pip install -r requirements-test.txt
-python3 -m pip install -r requirements-jax.txt
-python3 -c 'import jax; print(jax.__version__)'
+python -m pip install -r requirements-test.txt
 
-# Run bazel test command. Double test timeouts to avoid flakes.
-bazel test --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 \
-    --build_tests_only --test_output=errors \
-    --cache_test_results=no \
-    -- //...
+# Run tests using pytest.
+N_JOBS=$(grep -c ^processor /proc/cpuinfo)
+python -m pytest -n "${N_JOBS}" haiku
 
 # Test docs still build.
 cd docs/
 pip install -r requirements.txt
 # TODO(tomhennigan) Add `make doctest`.
 make html
-
-deactivate