diff --git a/.gitignore b/.gitignore
index 1c77aaa1..e90f5a60 100644
--- a/.gitignore
+++ b/.gitignore
@@ -25,3 +25,6 @@ docs/sg_execution_times.rst
 .idea
 .vscode
 
+# Local tests
+test_venv/
+optax-*.whl
diff --git a/pyproject.toml b/pyproject.toml
index 65cd61aa..6200e550 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -102,6 +102,22 @@ ignore = [
     "E703",  # allow semicolons (for jupyter notebooks)
 ]
 
+[tool.pytype]
+inputs = ["optax"]
+disable = ["import-error"]
+keep_going = true
+
+[tool.pytest.ini_options]
+    addopts = "--pyargs optax --numprocesses auto"
+
+[tool.pylint.main]
+ignore-patterns = [
+    ".*.pyi",
+]
+ignore = [
+    "test_venv",
+]
+
 [tool.pylint.messages_control]
 disable = [
     "bad-indentation",
diff --git a/test.sh b/test.sh
index bdb4fde6..b30213d8 100755
--- a/test.sh
+++ b/test.sh
@@ -13,27 +13,19 @@
 # limitations under the License.
 # ==============================================================================
 
-function cleanup {
-  deactivate
-  rm -r "${TEMP_DIR}"
-}
-trap cleanup EXIT
-
-REPO_DIR=$(pwd)
-TEMP_DIR=$(mktemp --directory)
-
-set -o errexit
 set -o nounset
-set -o pipefail
 
-# Install deps in a virtual env.
-python3 -m venv "${TEMP_DIR}/test_venv"
-source "${TEMP_DIR}/test_venv/bin/activate"
+echo "Creating virtual environment"
+python3 -m venv test_venv
+
+echo "Activating virtual environment"
+source test_venv/bin/activate
 
-# Install dependencies.
-python3 -m pip install --quiet --upgrade pip setuptools wheel
-python3 -m pip install --quiet --upgrade flake8 pytest-xdist pylint pylint-exit
-python3 -m pip install --quiet --editable ".[test, examples]"
+for dep in pip setuptools wheel flake8 pytest-xdist pylint pylint-exit pytype build ruff typing_extensions ".[test,examples,docs]"
+do
+  echo "Installing" $dep
+  python3 -m pip install --quiet --upgrade $dep
+done
 
 # Dp-accounting specifies exact minor versions as requirements which sometimes
 # become incompatible with other libraries optax needs. We therefore install
@@ -43,12 +35,15 @@ python3 -m pip install --quiet --editable ".[test, examples]"
 python3 -m pip install --quiet --editable ".[dp-accounting]"
 python3 -m pip install --quiet --no-deps "dp-accounting>=0.1.1"
 
-# Install the requested JAX version
-if [ -z "${JAX_VERSION-}" ]; then
+echo "Installing requested JAX version: ${JAX_VERSION-}"
+if [ -z "${JAX_VERSION-}" ]
+then
   : # use version installed in requirements above
-elif [ "$JAX_VERSION" = "newest" ]; then
+elif [ "$JAX_VERSION" = "newest" ]
+then
   python3 -m pip install --quiet --upgrade jax jaxlib
-elif [ "$JAX_VERSION" = "nightly" ]; then
+elif [ "$JAX_VERSION" = "nightly" ]
+then
   python3 -m pip install --quiet --upgrade --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
 else
   python3 -m pip install --quiet "jax==${JAX_VERSION}" "jaxlib==${JAX_VERSION}"
@@ -57,32 +52,31 @@ fi
 # Ensure optax was not installed by one of the dependencies above,
 # since if it is, the tests below will be run against that version instead of
 # the branch build.
+echo "Uninstalling optax (if installed)"
 python3 -m pip uninstall --quiet --yes optax
 
-# Lint with flake8.
-python3 -m flake8 --select=E9,F63,F7,F82,E225,E251 --show-source --statistics
+echo "Linting with flake8"
+flake8 --select=E9,F63,F7,F82,E225,E251 --show-source --statistics --exclude=build,dist,test_venv
 
-# Lint with pylint.
+echo "Linting with pylint"
 pylint .
 
-# Build the package.
-python3 -m pip install --quiet build
+echo "Building package"
 python3 -m build
-python3 -m pip wheel --no-deps dist/optax-*.tar.gz --wheel-dir "${TEMP_DIR}"
-python3 -m pip install --quiet "${TEMP_DIR}/optax-"*.whl
 
-# Check types with pytype.
-python3 -m pip install --quiet pytype
-pytype "optax" --keep-going --disable import-error
+echo "Building wheel"
+python3 -m pip wheel --no-deps dist/optax-*.tar.gz
 
-# Run tests using pytest.
-# Change directory to avoid importing the package from repo root.
-cd "${TEMP_DIR}"
-python3 -m pytest --numprocesses auto --pyargs optax
-cd "${REPO_DIR}"
+echo "Installing wheel"
+python3 -m pip install --quiet optax-*.whl
 
-# Build Sphinx docs.
-python3 -m pip install --quiet --editable ".[docs]"
+echo "Type checking with pytype"
+pytype
+
+echo "Running tests with pytest"
+pytest
+
+echo "Building Sphinx docs"
 # NOTE(vroulet) We have dependencies issues:
 # tensorflow > 2.13.1 requires ml-dtypes <= 0.3.2
 # but jax requires ml-dtypes >= 0.4.0
@@ -92,13 +86,14 @@ python3 -m pip install --quiet --editable ".[docs]"
 # bug (which issues conflict warnings but runs fine).
 # A long term solution is probably to fully remove tensorflow from our
 # dependencies.
-python3 -m pip install --upgrade --verbose typing_extensions
 cd docs
 make html
-make doctest # run doctests
+
+echo "Running doctests"
+make doctest
 cd ..
 
-pip install -U ruff
+echo "Running ruff"
 ruff check .
 
 echo "All tests passed. Congrats!"