Do not import scipy.signal and matplotlib at package import time

CHANGELOG.md

@@ -18,7 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
--
+- Stopped importing `scipy.signal` and `matplotlib` at `import torchmetrics` time, resolving both on first use instead ([#3463](https://github.com/Lightning-AI/torchmetrics/pull/3463))
 
 
 ### Deprecated

src/torchmetrics/__init__.py

@@ -27,13 +27,6 @@
     if not hasattr(PIL, "PILLOW_VERSION"):
         PIL.PILLOW_VERSION = PIL.__version__
 
-if package_available("scipy"):
-    import scipy.signal
-
-    # back compatibility patch due to SMRMpy using scipy.signal.hamming
-    if not hasattr(scipy.signal, "hamming"):
-        scipy.signal.hamming = scipy.signal.windows.hamming
-
 from torchmetrics import functional  # noqa: E402
 from torchmetrics.aggregation import (  # noqa: E402
     CatMetric,

src/torchmetrics/audio/__init__.py

@@ -29,17 +29,9 @@
     _PESQ_AVAILABLE,
     _PYSTOI_AVAILABLE,
     _REQUESTS_AVAILABLE,
-    _SCIPI_AVAILABLE,
     _TORCHAUDIO_AVAILABLE,
 )
 
-if _SCIPI_AVAILABLE:
-    import scipy.signal
-
-    # back compatibility patch due to SMRMpy using scipy.signal.hamming
-    if not hasattr(scipy.signal, "hamming"):
-        scipy.signal.hamming = scipy.signal.windows.hamming
-
 __all__ = [
     "ComplexScaleInvariantSignalNoiseRatio",
     "PermutationInvariantTraining",

src/torchmetrics/collections.py

@@ -26,7 +26,7 @@
 from torchmetrics.utilities import rank_zero_warn
 from torchmetrics.utilities.data import _flatten, _flatten_dict, allclose
 from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
-from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, _is_axes, plot_single_or_multi_val
 
 if not _MATPLOTLIB_AVAILABLE:
     __doctest_skip__ = ["MetricCollection.plot", "MetricCollection.plot_all"]
@@ -706,12 +706,12 @@ def plot(
         if not isinstance(together, bool):
             raise ValueError(f"Expected argument `together` to be a boolean, but got {type(together)}")
         if ax is not None:
-            if together and not isinstance(ax, _AX_TYPE):
+            if together and not _is_axes(ax):
                 raise ValueError(
                     f"Expected argument `ax` to be a matplotlib axis object, but got {type(ax)} when `together=True`"
                 )
             if not together and not (
-                isinstance(ax, Sequence) and all(isinstance(a, _AX_TYPE) for a in ax) and len(ax) == len(self)
+                isinstance(ax, Sequence) and all(_is_axes(a) for a in ax) and len(ax) == len(self)
             ):
                 raise ValueError(
                     f"Expected argument `ax` to be a sequence of matplotlib axis objects with the same length as the "

src/torchmetrics/functional/audio/__init__.py

@@ -29,17 +29,9 @@
     _PESQ_AVAILABLE,
     _PYSTOI_AVAILABLE,
     _REQUESTS_AVAILABLE,
-    _SCIPI_AVAILABLE,
     _TORCHAUDIO_AVAILABLE,
 )
 
-if _SCIPI_AVAILABLE:
-    import scipy.signal
-
-    # back compatibility patch due to SMRMpy using scipy.signal.hamming
-    if not hasattr(scipy.signal, "hamming"):
-        scipy.signal.hamming = scipy.signal.windows.hamming
-
 __all__ = [
     "complex_scale_invariant_signal_noise_ratio",
     "permutation_invariant_training",

src/torchmetrics/utilities/plot.py

@@ -12,45 +12,65 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections.abc import Generator, Sequence
+from contextlib import contextmanager
 from itertools import product
 from math import ceil, floor, sqrt
-from typing import Any, List, Optional, Union, no_type_check
+from typing import TYPE_CHECKING, Any, List, Optional, Union, no_type_check
 
 import numpy as np
 import torch
 from torch import Tensor
 
 from torchmetrics.utilities.imports import _LATEX_AVAILABLE, _MATPLOTLIB_AVAILABLE, _SCIENCEPLOT_AVAILABLE
 
-if _MATPLOTLIB_AVAILABLE:
+if TYPE_CHECKING:
     import matplotlib
     import matplotlib.axes
     import matplotlib.pyplot as plt
 
     _PLOT_OUT_TYPE = tuple[plt.Figure, Union[matplotlib.axes.Axes, np.ndarray]]
     _AX_TYPE = matplotlib.axes.Axes
     _CMAP_TYPE = Union[matplotlib.colors.Colormap, str]
-
-    style_change = plt.style.context
 else:
-    _PLOT_OUT_TYPE = tuple[object, object]  # type: ignore[misc]
+    # `matplotlib` is resolved on first use instead of at import: pulling it in eagerly costs roughly a second of
+    # `import torchmetrics` for everyone, including those who never plot. The ~140 modules that annotate with these
+    # aliases never introspect the annotations at runtime, so plain `object` is enough outside of type checking.
+    _PLOT_OUT_TYPE = tuple[object, object]
     _AX_TYPE = object
-    _CMAP_TYPE = object  # type: ignore[misc]
+    _CMAP_TYPE = object
+
+_style = ["science"] if _SCIENCEPLOT_AVAILABLE and _LATEX_AVAILABLE else ["default"]
 
-    from contextlib import contextmanager
 
-    @contextmanager
-    def style_change(*args: Any, **kwargs: Any) -> Generator:
-        """No-ops decorator if matplotlib is not installed."""
+@contextmanager
+def style_change(*args: Any, **kwargs: Any) -> Generator:
+    """Apply a ``matplotlib`` style, no-op if ``matplotlib`` is not installed."""
+    if not _MATPLOTLIB_AVAILABLE:
         yield
+        return
 
+    import matplotlib.pyplot as plt
 
-if _SCIENCEPLOT_AVAILABLE:
-    import scienceplots  # noqa: F401
+    if _SCIENCEPLOT_AVAILABLE:
+        import scienceplots  # noqa: F401  # registers the "science" style on import
 
-    _style = ["science", "no-latex"]
+    with plt.style.context(*args, **kwargs):
+        yield
 
-_style = ["science"] if _SCIENCEPLOT_AVAILABLE and _LATEX_AVAILABLE else ["default"]
+
+def _is_axes(obj: Any) -> bool:
+    """Check whether ``obj`` is a ``matplotlib`` ``Axes``.
+
+    The ``_AX_TYPE`` alias is only a real class while type checking, so it cannot be used with ``isinstance``. This
+    resolves the actual class instead, importing ``matplotlib`` only when there is something plausible to check.
+
+    """
+    if not _MATPLOTLIB_AVAILABLE:
+        return False
+
+    import matplotlib.axes
+
+    return isinstance(obj, matplotlib.axes.Axes)
 
 
 def _error_on_missing_matplotlib() -> None:
@@ -93,6 +113,8 @@ def plot_single_or_multi_val(
 
     """
     _error_on_missing_matplotlib()
+    import matplotlib.pyplot as plt
+
     fig, ax = plt.subplots() if ax is None else (None, ax)
     ax.get_xaxis().set_visible(False)
 
@@ -201,13 +223,13 @@ def _get_text_color(patch_color: tuple[float, float, float, float]) -> str:
     return ".1" if y > 0.4 else "white"
 
 
-def trim_axs(axs: Union[_AX_TYPE, np.ndarray], nb: int) -> Union[np.ndarray, _AX_TYPE]:  # type: ignore[valid-type]
+def trim_axs(axs: Union[_AX_TYPE, np.ndarray], nb: int) -> Union[np.ndarray, _AX_TYPE]:
     """Reduce `axs` to `nb` Axes.
 
     All further Axes are removed from the figure.
 
     """
-    if isinstance(axs, _AX_TYPE):
+    if _is_axes(axs):
         return axs
 
     axs = axs.flat  # type: ignore[union-attr]
@@ -248,6 +270,7 @@ def plot_confusion_matrix(
 
     """
     _error_on_missing_matplotlib()
+    import matplotlib.pyplot as plt
 
     if confmat.ndim == 3:  # multilabel
         nb, n_classes = confmat.shape[0], 2
@@ -297,7 +320,7 @@ def plot_confusion_matrix(
 def plot_curve(
     curve: Union[tuple[Tensor, Tensor, Tensor], tuple[List[Tensor], List[Tensor], List[Tensor]]],
     score: Optional[Tensor] = None,
-    ax: Optional[_AX_TYPE] = None,  # type: ignore[valid-type]
+    ax: Optional[_AX_TYPE] = None,
     label_names: Optional[tuple[str, str]] = None,
     legend_name: Optional[str] = None,
     name: Optional[str] = None,
@@ -331,6 +354,8 @@ def plot_curve(
     x, y = curve[:2]
 
     _error_on_missing_matplotlib()
+    import matplotlib.pyplot as plt
+
     fig, ax = plt.subplots() if ax is None else (None, ax)
 
     if isinstance(x, Tensor) and isinstance(y, Tensor) and x.ndim == 1 and y.ndim == 1:

src/torchmetrics/wrappers/multitask.py

@@ -21,7 +21,7 @@
 from torchmetrics.collections import MetricCollection
 from torchmetrics.metric import Metric
 from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
-from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
+from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, _is_axes
 from torchmetrics.wrappers.abstract import WrapperMetric
 
 if not _MATPLOTLIB_AVAILABLE:
@@ -358,7 +358,7 @@ def plot(
             if not isinstance(axes, Sequence):
                 raise TypeError(f"Expected argument `axes` to be a Sequence. Found type(axes) = {type(axes)}")
 
-            if not all(isinstance(ax, _AX_TYPE) for ax in axes):
+            if not all(_is_axes(ax) for ax in axes):
                 raise TypeError("Expected each ax in argument `axes` to be a matplotlib axis object")
 
             if len(axes) != len(self.task_metrics):

tests/unittests/audio/test_srmr.py

@@ -15,10 +15,17 @@
 from typing import Any
 
 import pytest
+import scipy.signal
 import torch
-from srmrpy import srmr as srmrpy_srmr
 from torch import Tensor
 
+# back compatibility patch due to SRMRpy using `scipy.signal.hamming`, which was removed in scipy 1.13. Applied here
+# rather than in `torchmetrics/__init__.py` so that importing torchmetrics does not pull in scipy, see #3457.
+if not hasattr(scipy.signal, "hamming"):
+    scipy.signal.hamming = scipy.signal.windows.hamming
+
+from srmrpy import srmr as srmrpy_srmr
+
 from torchmetrics.audio.srmr import SpeechReverberationModulationEnergyRatio
 from torchmetrics.functional.audio.srmr import speech_reverberation_modulation_energy_ratio
 from unittests._helpers import seed_all

tests/unittests/utilities/test_lazy_imports.py

@@ -0,0 +1,121 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Optional heavy dependencies must not be imported just because they happen to be installed.
+
+`import torchmetrics` is on the critical path for downstream libraries — Lightning imports it merely to compare versions
+— so eagerly pulling in plotting or signal-processing stacks costs every user, including those who never touch the
+features that need them. See
+https://github.com/Lightning-AI/torchmetrics/issues/3457.
+
+"""
+
+import subprocess
+import sys
+
+import pytest
+
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE, _SCIPI_AVAILABLE
+
+# checked in a subprocess: the test session itself has already imported large parts of the world
+_PROBE = (
+    "import sys, torchmetrics; "
+    "print(','.join(m for m in ('scipy.signal', 'matplotlib', 'matplotlib.pyplot') if m in sys.modules))"
+)
+
+
+def _modules_after_importing_torchmetrics() -> set:
+    """Return the probed optional modules that a bare ``import torchmetrics`` pulled in."""
+    # S603 is suppressed below: the interpreter path and the probe are fixed literals, no external input is involved
+    out = subprocess.run(  # noqa: S603
+        [sys.executable, "-c", _PROBE], capture_output=True, text=True, check=True
+    )
+    return {name for name in out.stdout.strip().split(",") if name}
+
+
+@pytest.mark.skipif(not _SCIPI_AVAILABLE, reason="test only meaningful when scipy is installed")
+def test_importing_torchmetrics_does_not_import_scipy_signal():
+    """`scipy.signal` was only pulled in by the SRMRpy back-compat patch, which no longer runs at import.
+
+    Asserted on `scipy.signal` rather than `scipy` on purpose: when `transformers` is installed,
+    `functional/text/bert.py` imports it at module level and that pulls in `scipy.sparse`. That is a separate eager
+    import from the one this test guards, so asserting on the `scipy` root would make this test fail for an unrelated
+    reason.
+
+    """
+    loaded = _modules_after_importing_torchmetrics()
+    assert "scipy.signal" not in loaded, f"`import torchmetrics` pulled in {sorted(loaded)}"
+
+
+@pytest.mark.skipif(not _MATPLOTLIB_AVAILABLE, reason="test only meaningful when matplotlib is installed")
+def test_importing_torchmetrics_does_not_import_matplotlib():
+    """`matplotlib` is only needed by ``.plot()``, which most users never call."""
+    loaded = _modules_after_importing_torchmetrics()
+    assert "matplotlib" not in loaded, f"`import torchmetrics` pulled in {sorted(loaded)}"
+    assert "matplotlib.pyplot" not in loaded
+
+
+@pytest.mark.skipif(not _MATPLOTLIB_AVAILABLE, reason="requires matplotlib")
+def test_plot_type_aliases_are_defined_at_runtime():
+    """The aliases are annotations on ~140 metrics, so they must still resolve without matplotlib loaded."""
+    from torchmetrics.utilities.plot import _AX_TYPE, _CMAP_TYPE, _PLOT_OUT_TYPE
+
+    assert _AX_TYPE is not None
+    assert _CMAP_TYPE is not None
+    assert _PLOT_OUT_TYPE is not None
+
+
+@pytest.mark.skipif(not _MATPLOTLIB_AVAILABLE, reason="requires matplotlib")
+def test_is_axes_discriminates_real_axes():
+    """``_AX_TYPE`` is ``object`` at runtime, so ``isinstance(x, _AX_TYPE)`` would match anything.
+
+    Several call sites branch on whether an argument is a real ``Axes`` — ``trim_axs``, ``MetricCollection.plot`` and
+    ``MultitaskWrapper.plot``. They must use ``_is_axes``, otherwise the checks silently pass for every input:
+    validation stops raising and ``trim_axs`` returns the untrimmed array.
+
+    """
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    from torchmetrics.utilities.plot import _AX_TYPE, _is_axes
+
+    _, ax = plt.subplots()
+    _, axs = plt.subplots(nrows=2, ncols=2)
+
+    assert _is_axes(ax)
+    assert not _is_axes(axs)  # an ndarray of axes, not an Axes
+    assert not _is_axes("not an axis")
+    assert not _is_axes(None)
+
+    # the alias itself must not be used for this: every object is an instance of `object`
+    assert isinstance("not an axis", _AX_TYPE)
+
+    plt.close("all")
+
+
+def test_style_change_works_as_context_manager_and_decorator():
+    """``style_change`` is applied as a decorator at import time, so it must not need matplotlib to be constructed."""
+    from torchmetrics.utilities.plot import _style, style_change
+
+    with style_change(_style):
+        pass
+
+    @style_change(_style)
+    def _fn() -> str:
+        return "ok"
+
+    # a context manager used as a decorator must survive being called more than once
+    assert _fn() == "ok"
+    assert _fn() == "ok"