Install virtual environment in current directory instead of /tmp

This solves the issue of pytest 8.0.0 not being able to run
tests from /tmp/ directory.

See pytest-dev/pytest#11807

test.sh

@@ -17,8 +17,9 @@
 set -xeuo pipefail
 
 # Install deps in a virtual env.
-readonly VENV_DIR=/tmp/optax-env
-rm -rf "${VENV_DIR}"
+mkdir -p _testing
+readonly VENV_DIR="$(mktemp -d -p `pwd`/_testing optax-env.XXXXXXXX)"
+# in the unlikely case in which there was something in that directory
 python3 -m venv "${VENV_DIR}"
 source "${VENV_DIR}/bin/activate"
 python --version
@@ -65,11 +66,11 @@ pip install optax*.whl
 
 # Check types with pytype.
 pip install pytype
-pytype `find optax/_src/ examples optax/contrib -name '*.py' | xargs` -k -d import-error
+pytype `find optax/_src examples optax/contrib -name '*.py' | xargs` -k -d import-error
 
 # Run tests using pytest.
 # Change directory to avoid importing the package from repo root.
-mkdir _testing && cd _testing
+cd _testing
 python -m pytest -n auto --pyargs optax
 cd ..