Repository: google/flax

Issue:
Potentially make matplotlib an optional dependency
Apologies in advance for quoting, nearly verbatim, these issues from the [ott-jax](https://github.com/ott-jax/ott/issues/292) and [jaxopt](https://github.com/google/jaxopt/issues/396) packages but the issue is largely identical.

matplotlib is a fairly "heavy" dependency in that it pulls in lots of extra packages. However, it's not required to use the bulk of `flax` outside of potentially parts of [flax/metrics/tensorboard.py](https://github.com/google/flax/blob/1e7d843909faf48c034ea12eb7bf3d007efe4f7b/flax/metrics/tensorboard.py).

I'm wondering whether the team would consider either removing it or making it an optional dependency to minimize the number of packages pulled in by downstream projects.


Discussion:
I agree with this, we can create an e.g. `all` group and bundle `matplotlib` there:

```bash
pip install flax[all]
```

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -25,38 +25,40 @@
   README = ""
 
 install_requires = [
-    "numpy>=1.12",
-    "jax>=0.4.2",
-    "matplotlib",  # only needed for tensorboard export
-    "msgpack",
-    "optax",
-    "orbax",
-    "tensorstore",
-    "rich>=11.1",
-    "typing_extensions>=4.1.1",
-    "PyYAML>=5.4.1",
+  "numpy>=1.12",
+  "jax>=0.4.2",
+  "msgpack",
+  "optax",
+  "orbax",
+  "tensorstore",
+  "rich>=11.1",
+  "typing_extensions>=4.1.1",
+  "PyYAML>=5.4.1",
+]
+all_requires = [
+  "matplotlib",  # only needed for tensorboard export
 ]
 
 tests_require = [
-    "atari-py==0.2.5",  # Last version does not have the ROMs we test on pre-packaged
-    "clu",  # All examples.
-    "gym==0.18.3",
-    "jaxlib",
-    "jraph>=0.0.6dev0",
-    "ml-collections",
-    "mypy",
-    "opencv-python",
-    "pytest",
-    "pytest-cov",
-    "pytest-custom_exit_code",
-    "pytest-xdist==1.34.0",  # upgrading to 2.0 broke tests, need to investigate
-    "pytype",
-    "sentencepiece",  # WMT/LM1B examples
-    "tensorflow_text>=2.11.0",  # WMT/LM1B examples
-    "tensorflow_datasets",
-    "tensorflow",
-    "torch",
-    "nbstripout",
+  "atari-py==0.2.5",  # Last version does not have the ROMs we test on pre-packaged
+  "clu",  # All examples.
+  "gym==0.18.3",
+  "jaxlib",
+  "jraph>=0.0.6dev0",
+  "ml-collections",
+  "mypy",
+  "opencv-python",
+  "pytest",
+  "pytest-cov",
+  "pytest-custom_exit_code",
+  "pytest-xdist==1.34.0",  # upgrading to 2.0 broke tests, need to investigate
+  "pytype",
+  "sentencepiece",  # WMT/LM1B examples
+  "tensorflow_text>=2.11.0",  # WMT/LM1B examples
+  "tensorflow_datasets",
+  "tensorflow",
+  "torch",
+  "nbstripout",
 ]
 
 __version__ = None
@@ -65,28 +67,29 @@
   exec(f.read(), globals())
 
 setup(
-    name="flax",
-    version=__version__,
-    description="Flax: A neural network library for JAX designed for flexibility",
-    long_description="\n\n".join([README]),
-    long_description_content_type="text/markdown",
-    classifiers=[
-        "Development Status :: 3 - Alpha",
-        "Intended Audience :: Developers",
-        "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: Apache Software License",
-        "Programming Language :: Python :: 3.7",
-        "Topic :: Scientific/Engineering :: Artificial Intelligence",
-        ],
-    keywords="",
-    author="Flax team",
-    author_email="flax-dev@google.com",
-    url="https://github.com/google/flax",
-    packages=find_packages(),
-    package_data={"flax": ["py.typed"]},
-    zip_safe=False,
-    install_requires=install_requires,
-    extras_require={
-        "testing": tests_require,
-        },
-    )
+  name="flax",
+  version=__version__,
+  description="Flax: A neural network library for JAX designed for flexibility",
+  long_description="\n\n".join([README]),
+  long_description_content_type="text/markdown",
+  classifiers=[
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3.7",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  ],
+  keywords="",
+  author="Flax team",
+  author_email="flax-dev@google.com",
+  url="https://github.com/google/flax",
+  packages=find_packages(),
+  package_data={"flax": ["py.typed"]},
+  zip_safe=False,
+  install_requires=install_requires,
+  extras_require={
+    "testing": tests_require,
+    "all": all_requires,
+  },
+)


Test patch: