Clean dependencies

PiperOrigin-RevId: 748836033

Author: vroulet

docs/conf.py

@@ -199,7 +199,9 @@ def _recursive_add_annotations_import():
     'cifar10_resnet.ipynb',
     'adversarial_training.ipynb',
     'reduce_on_plateau.ipynb',
-    'differentially_private_sgd.ipynb'
+    'differentially_private_sgd.ipynb',
+    'lookahead_mnist.ipynb',
+    'mlp_mnist.ipynb',
 ]
 
 # -- Options for katex ------------------------------------------------------

examples/adversarial_training.ipynb

@@ -11,7 +11,6 @@
         "\n",
         "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/optax/blob/main/examples/adversarial_training.ipynb)\n",
         "\n",
-        "\n",
         "The following code trains a convolutional neural network (CNN) to be robust\n",
         "with respect to the projected gradient descent (PGD) method.\n",
         "\n",
@@ -28,7 +27,11 @@
         "  and harnessing adversarial examples.\", https://arxiv.org/abs/1412.6572\n",
         "\n",
         "  Madry, Aleksander, et al. \"Towards deep learning models resistant to\n",
-        "  adversarial attacks.\", https://arxiv.org/abs/1706.06083"
+        "  adversarial attacks.\", https://arxiv.org/abs/1706.06083\n",
+        "\n",
+        "\n",
+        "To run the colab locally you need install the\n",
+        "`tensorflow`, `tensorflow-datasets` packages via `pip`."
       ]
     },
     {

examples/cifar10_resnet.ipynb

@@ -10,7 +10,10 @@
         "\n",
         "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/google-deepmind/optax/blob/main/examples/cifar10_resnet.ipynb)\n",
         "\n",
-        "This notebook trains a residual network (ResNet) with optax on CIFAR10 or CIFAR100."
+        "This notebook trains a residual network (ResNet) with optax on CIFAR10 or CIFAR100.\n",
+        "\n",
+        "To run the colab locally you need install the\n",
+        "`tensorflow`, `tensorflow-datasets` packages via `pip`."
       ]
     },
     {

examples/contrib/differentially_private_sgd.ipynb

@@ -14,6 +14,10 @@
         "example in the [JAX repo](\n",
         "https://github.com/jax-ml/jax/blob/main/examples/differentially_private_sgd.py).\n",
         "\n",
+        "To run the colab locally you need install the\n",
+        "`dp-accounting`, `tensorflow`, `tensorflow-datasets`, packages via `pip`.\n",
+        "\n",
+        "\n",
         "[Differentially Private Stochastic Gradient Descent](https://arxiv.org/abs/1607.00133) requires clipping the per-example parameter\n",
         "gradients, which is non-trivial to implement efficiently for convolutional\n",
         "neural networks.  The JAX XLA compiler shines in this setting by optimizing the\n",

examples/contrib/reduce_on_plateau.ipynb

@@ -10,7 +10,10 @@
         "\n",
         "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/google-deepmind/optax/blob/main/examples/contrib/reduce_on_plateau.ipynb)\n",
         "\n",
-        "In this notebook, we explore the power of {py:func}`optax.contrib.reduce_on_plateau` scheduler, that reduces the learning rate when a metric has stopped improving. We will be solving a classification task by training a simple Multilayer Perceptron (MLP) on the fashion MNIST dataset."
+        "In this notebook, we explore the power of {py:func}`optax.contrib.reduce_on_plateau` scheduler, that reduces the learning rate when a metric has stopped improving. We will be solving a classification task by training a simple Multilayer Perceptron (MLP) on the fashion MNIST dataset.\n",
+        "\n",
+        "To run the colab locally you need install the\n",
+        "`tensorflow`, `tensorflow-datasets` packages via `pip`."
       ]
     },
     {

examples/lookahead_mnist.ipynb

@@ -10,7 +10,10 @@
         "\n",
         "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/google-deepmind/optax/blob/main/examples/lookahead_mnist.ipynb)\n",
         "\n",
-        "This notebook trains a simple Convolution Neural Network (CNN) for hand-written digit recognition (MNIST dataset) using {py:func}`optax.lookahead`."
+        "This notebook trains a simple Convolution Neural Network (CNN) for hand-written digit recognition (MNIST dataset) using {py:func}`optax.lookahead`.\n",
+        "\n",
+        "To run the colab locally you need install the\n",
+        "`tensorflow`, `tensorflow-datasets` packages via `pip`."
       ]
     },
     {

examples/mlp_mnist.ipynb

@@ -10,7 +10,10 @@
         "\n",
         "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.sandbox.google.com/github/google-deepmind/optax/blob/main/examples/mlp_mnist.ipynb)\n",
         "\n",
-        "This notebook trains a simple Multilayer Perceptron (MLP) classifier for hand-written digit recognition (MNIST dataset)."
+        "This notebook trains a simple Multilayer Perceptron (MLP) classifier for hand-written digit recognition (MNIST dataset).\n",
+        "\n",
+        "To run the colab locally you need install the\n",
+        "`tensorflow`, `tensorflow-datasets` packages via `pip`."
       ]
     },
     {

pyproject.toml

@@ -45,20 +45,11 @@ documentation = "https://optax.readthedocs.io/"
 
 [project.optional-dependencies]
 test = [
-    "dm-tree>=0.1.7",
     "flax>=0.5.3",
     "scipy>=1.7.1",
     "scikit-learn"
 ]
 
-examples = [
-  "tensorflow-datasets>=4.2.0",
-  "tensorflow>=2.4.0",
-  "dp_accounting>=0.4",
-  "ipywidgets",
-  "flax",
-]
-
 docs = [
     "sphinx>=6.0.0",
     "sphinx-book-theme>=1.0.1",  # Older versions fail to pin pydata-sphinx-theme
@@ -69,20 +60,10 @@ docs = [
     "matplotlib>=3.5.0",
     "sphinx-gallery>=0.14.0",
     "sphinx-collections>=0.0.1",
-    "tensorflow>=2.4.0",
-    "tensorflow-datasets>=4.2.0",
     "flax",
     "sphinx_contributors",
 ]
 
-dp-accounting = [
-    "absl-py>=1.0.0",
-    "attrs>=21.4.0",
-    "mpmath>=1.2.1",
-    "numpy>=1.21.4",
-    "scipy>=1.7.1"
-]
-
 [tool.setuptools.packages.find]
 include = ["README.md", "LICENSE"]
 exclude = ["*_test.py"]

test.sh

@@ -38,15 +38,7 @@ pre-commit run -a
 # Install dependencies.
 python3 -m uv pip install --quiet --upgrade pip setuptools wheel
 python3 -m uv pip install --quiet --upgrade flake8 pytest-xdist pylint pylint-exit
-python3 -m uv pip install --quiet --editable ".[test, examples]"
-
-# Dp-accounting specifies exact minor versions as requirements which sometimes
-# become incompatible with other libraries optax needs. We therefore install
-# dependencies for dp-accounting manually.
-# TODO(b/239416992): Remove this workaround if dp-accounting switches to minimum
-# version requirements.
-python3 -m uv pip install --quiet --editable ".[dp-accounting]"
-python3 -m uv pip install --quiet --no-deps "dp-accounting>=0.1.1"
+python3 -m uv pip install --quiet --editable ".[test]"
 
 # Install the requested JAX version
 if [ -z "${JAX_VERSION-}" ]; then
@@ -89,16 +81,6 @@ cd "${REPO_DIR}"
 
 # Build Sphinx docs.
 python3 -m uv pip install --quiet --editable ".[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
-# So the environment is solved with tensorflow == 2.13.1 which requires
-# typing_extensions < 4.6, which in turn prevents the import of TypeAliasType in
-# IPython. We solve it here by simply upgrading typing_extensions to avoid that
-# bug (which issues conflict warnings but runs fine).
-# A long term solution is probably to fully remove tensorflow from our
-# dependencies.
-python3 -m uv pip install --upgrade --verbose typing_extensions
 cd docs
 make html
 make doctest # run doctests