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

[VenishPaneliya]

What does this PR do?

Refs #3457 — the scipy.signal and matplotlib halves. Leaves the torchvision/arniqa half to #3432 to avoid conflicting.

import torchmetrics pulled in two optional stacks that nothing on the import path uses. This sits on the critical path for downstream libraries — Lightning imports torchmetrics only to compare versions — so everyone paid for it, including users who never plot or touch audio metrics.

scipy.signal

Came from the SRMRpy hamming back-compat patch, duplicated across three eagerly executed __init__.py files. srmrpy is only ever imported by tests/unittests/audio/test_srmr.py — the library reimplements SRMR in pure PyTorch and never imports it — so the patch moves into that test module.

matplotlib

Came from utilities/plot.py, which imported it at module level purely to build the _AX_TYPE / _PLOT_OUT_TYPE / _CMAP_TYPE annotation aliases and to alias style_change = plt.style.context. The aliases now resolve to the real classes under TYPE_CHECKING and to object at runtime, and style_change imports pyplot on first use.

I kept runtime object fallbacks rather than making the aliases TYPE_CHECKING-only, because that would require from __future__ import annotations in the ~140 modules that annotate with them. This way the diff stays inside plot.py (plus the two call sites below).

The subtle part

Once the aliases stop being real classes at runtime, isinstance(x, _AX_TYPE) becomes isinstance(x, object) — always true. Four places relied on that check:

Location	What silently broke
plot.py trim_axs	returned the untrimmed axes array instead of flattening and trimming
collections.py ×2	MetricCollection.plot argument validation stopped raising
wrappers/multitask.py	MultitaskWrapper.plot argument validation stopped raising

The trim_axs one was caught by test_confusion_matrix_plotter[multilabel confusion matrix] failing with AttributeError: 'numpy.ndarray' object has no attribute 'set_title'. All four now go through a new _is_axes helper that resolves the actual class, and there is a regression test pinning the behaviour.

Measurements

import torchmetrics, median of 5 cold subprocesses, Python 3.13 / Windows:

With only scipy and matplotlib installed — the configuration this PR targets:

	master	this PR
wall time	5559 ms	3571 ms (−36%)
sys.modules	2158	1371 (−787)
heavy modules loaded	scipy.signal, matplotlib.pyplot	none

With transformers and torchvision also installed, where other eager imports dominate:

	master	this PR
wall time	15686 ms	14426 ms (−8%)
sys.modules	3947	3776 (−171)
heavy modules loaded	scipy.signal, matplotlib.pyplot, torchvision	torchvision

Tests

New tests/unittests/utilities/test_lazy_imports.py checks in a subprocess that a bare import torchmetrics leaves scipy.signal and matplotlib out of sys.modules, that the aliases still resolve at runtime, that style_change works as both a context manager and a repeatable decorator, and that _is_axes discriminates where the alias cannot.

$ MPLBACKEND=Agg pytest tests/unittests/utilities -q
6 failed, 371 passed, 8 skipped in 171.62s

All 6 failures are missing optional packages locally (pesq, gammatone, pycocotools) and are identical on unpatched master:

$ git stash && MPLBACKEND=Agg pytest tests/unittests/utilities -q
9 failed, 368 passed, 8 skipped in 110.95s

— the 3 extra being this PR's own new tests, which correctly fail without the fix.

Lint and types, compared against master for the same files:

$ ruff check <changed files>
3 errors  (RUF059 ×3 — all pre-existing, identical on master)

$ ruff format --check <changed files>
5 files already formatted

$ mypy src/torchmetrics/utilities/plot.py src/torchmetrics/collections.py src/torchmetrics/wrappers/multitask.py
Found 11 errors   (master: 16 — the TYPE_CHECKING types let mypy resolve more, net -5)

One thing found along the way, not fixed here

When transformers is installed, import torchmetrics still pulls in scipy — but via a different path than this issue describes:

torchmetrics/__init__.py
  → functional/__init__.py:129
    → functional/text/__init__.py:50
      → functional/text/bert.py:56  →  from transformers import AutoModel, AutoTokenizer

bert.py imports transformers at module level, which pulls in scipy.sparse (and a lot more). That is a separate eager import from the SRMRpy one, so I left it alone and scoped the regression test to scipy.signal accordingly — asserting on the scipy root would make the test fail for an unrelated reason. Happy to open a follow-up issue if that is worth tracking.

Before submitting

• Was this discussed/agreed via a GitHub issue? (Import time: scipy.signal, torchvision and matplotlib are imported eagerly by import torchmetrics #3457)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs? (docstrings on the new helpers)
• Did you write any new necessary tests?

[VenishPaneliya]

Apologies — I should have checked the open PR list before opening this. #3459 already covers the scipy.signal half of #3457, and #3432 covers the torchvision half. That's my mistake.

Flagging the overlap so you can decide rather than review both:

• The scipy.signal change here is the same idea as Do not import scipy.signal at torchmetrics import time #3459 — drop the SRMRpy hamming patch from the three eager __init__.py files and move it into tests/unittests/audio/test_srmr.py, which is the only place srmrpy is ever imported.
• What isn't covered elsewhere is the matplotlib half: utilities/plot.py imported matplotlib at module level purely to build the _AX_TYPE / _PLOT_OUT_TYPE / _CMAP_TYPE aliases and to alias plt.style.context. That's the larger share of the win here (5559 ms -> 3571 ms, 2158 -> 1371 modules, measured with only scipy and matplotlib installed).
• One thing worth keeping regardless of which PR you take: once those aliases stop being real classes at runtime, isinstance(x, _AX_TYPE) becomes isinstance(x, object) and is always true. Four call sites rely on that check — trim_axs in plot.py, two in MetricCollection.plot, one in MultitaskWrapper.plot — so they silently stop trimming and stop validating, with no error raised. test_confusion_matrix_plotter[multilabel confusion matrix] catches it. This PR routes them through a small _is_axes helper and adds a regression test.

If you'd prefer, I'm happy to close this and rebase just the matplotlib half on top of #3459 once that lands, so the two don't conflict. Let me know which you'd rather have.