Merge pull request #746 from fabianp:pytest-tmp

PiperOrigin-RevId: 602343107

Author: OptaxDev

test.sh

@@ -17,8 +17,10 @@
 set -xeuo pipefail
 
 # Install deps in a virtual env.
-readonly VENV_DIR=/tmp/optax-env
-rm -rf "${VENV_DIR}"
+rm -rf _testing
+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 +67,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 ..