Make torch and the sklearn stack optional extras

.github/workflows/docs.yml

@@ -36,12 +36,17 @@ jobs:
           python-version: "3.12"
 
       - name: Install the project
+        # The docs group carries `ezmsg-learn[all]`: autodoc imports every
+        # documented module, including the ones behind the torch/sklearn extras.
         run: uv sync --only-group docs
 
       - name: Build documentation
+        # --no-sync so the build uses exactly the environment installed above.
+        # Without it, `uv run` re-syncs to the default groups, which quietly
+        # decides the docs dependencies for us.
         run: |
           cd docs
-          uv run make html
+          uv run --no-sync make html
 
       - name: Add .nojekyll file
         run: touch docs/build/html/.nojekyll

.github/workflows/python-tests.yml

@@ -47,3 +47,34 @@ jobs:
         run: uv run pytest tests/integration -v --tb=long
         env:
           PYTHONFAULTHANDLER: 1
+
+  # The backends are extras; this job proves a base install stays free of them
+  # and that the modules behind an extra say which extra they need.
+  minimal-install:
+    name: Minimal install (no extras)
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          python-version: "3.13"
+
+      - name: Install the project without extras
+        # Deliberately not `uv sync`: the dev dependency-groups pull in [all].
+        run: |
+          uv venv
+          uv pip install . pytest
+
+      - name: Assert no backend is installed
+        run: |
+          .venv/bin/python -c "
+          import importlib.util
+          present = [m for m in ('torch', 'sklearn', 'river', 'pandas') if importlib.util.find_spec(m)]
+          assert not present, f'base install pulled in {present}'
+          "
+
+      - name: Run the optional-dependency guards
+        run: .venv/bin/python -m pytest tests/unit/test_optional_deps.py -v

README.md

@@ -4,31 +4,39 @@ This repository contains a Python package with modules for machine learning (ML)
 
 > If you are only interested in offline analysis without concern for reproducibility in online applications, then you should probably look elsewhere.
 
-Processing units include dimensionality reduction, linear regression, and classification that can be initialized with known weights, or adapted on-the-fly with incoming (labeled) data. Machine-learning code depends on `river`, `scikit-learn`, `numpy`, and `torch`.
+Processing units include dimensionality reduction, linear regression, and classification that can be initialized with known weights, or adapted on-the-fly with incoming (labeled) data.
 
 ## Installation
 
-Install from PyPI:
+The base install is NumPy-only. The machine-learning backends are optional extras, so a deployment that uses only the lightweight processors does not pay for a PyTorch or scikit-learn install:
 
 ```bash
-pip install ezmsg-learn
+pip install ezmsg-learn              # numpy-only processors
+pip install "ezmsg-learn[sklearn]"   # + pandas, river, scikit-learn
+pip install "ezmsg-learn[torch]"     # + torch
+pip install "ezmsg-learn[all]"       # everything
 ```
 
 Or install the latest development version:
 
 ```bash
-pip install git+https://github.com/ezmsg-org/ezmsg-learn@dev
+pip install "git+https://github.com/ezmsg-org/ezmsg-learn@dev#egg=ezmsg-learn[all]"
 ```
 
+Importing a module whose backend is not installed raises an `ImportError` naming the extra to install.
+
 ## Dependencies
 
-- `ezmsg`
-- `ezmsg-baseproc`
-- `ezmsg-sigproc`
-- `numpy`
-- `scipy`
-- `scikit-learn`
-- `river`
+Base (`pip install ezmsg-learn`) — `ezmsg`, `ezmsg-baseproc`, `ezmsg-sigproc`, `numpy`, `scipy`, `array-api-compat`.
+
+| Extra | Adds | Covers |
+| --- | --- | --- |
+| _(none)_ | — | `process.ssr`, `process.flatten`, `process.seqseqsampler`, `process.refit_kalman`, `model.cca`, `model.refit_kalman` |
+| `sklearn` | `pandas`, `river`, `scikit-learn` | `process.adaptive_linear_regressor`, `process.linear_regressor`, `process.sgd`, `process.slda`, `process.sklearn`, `dim_reduce.*` |
+| `torch` | `torch` | `process.base`, `process.torch`, `process.rnn`, `process.transformer`, `process.mlp_old`, `model.mlp`, `model.rnn`, `model.transformer` |
+| `all` | both of the above | everything, including all `collection.sample_adapt_regressor` backends |
+
+`collection.sample_adapt_regressor` imports its backend lazily, so it needs only the extra for the `model_type` in use — and none at all for `model_type="kalman"`.
 
 
 ## Development

docs/source/index.rst

@@ -26,23 +26,68 @@ Most modules support both:
 Installation
 ------------
 
-Install directly from GitHub:
+The base install is NumPy-only; the machine-learning backends are optional
+extras, so a deployment that uses only the lightweight processors does not pay
+for a PyTorch or scikit-learn install:
 
 .. code-block:: bash
 
-   pip install git+https://github.com/ezmsg-org/ezmsg-learn
+   pip install ezmsg-learn              # numpy-only processors
+   pip install "ezmsg-learn[sklearn]"   # + pandas, river, scikit-learn
+   pip install "ezmsg-learn[torch]"     # + torch
+   pip install "ezmsg-learn[all]"       # everything
+
+Or install directly from GitHub:
+
+.. code-block:: bash
+
+   pip install "git+https://github.com/ezmsg-org/ezmsg-learn#egg=ezmsg-learn[all]"
+
+Importing a module whose backend is not installed raises an ``ImportError``
+naming the extra to install.
 
 Dependencies
 ^^^^^^^^^^^^
 
-This package requires:
+The base install requires:
 
 * ``ezmsg`` - Core ezmsg framework
+* ``ezmsg-baseproc`` - Processor base classes
 * ``ezmsg-sigproc`` - Signal processing extensions
 * ``numpy`` - Numerical computing
-* ``scikit-learn`` - Machine learning utilities
-* ``torch`` - Deep learning framework
-* ``river`` - Online machine learning
+* ``scipy`` - Scientific computing
+* ``array-api-compat`` - Array API portability layer
+
+Optional extras
+^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :widths: 15 25 60
+
+   * - Extra
+     - Adds
+     - Covers
+   * - *(none)*
+     - —
+     - ``process.ssr``, ``process.flatten``, ``process.seqseqsampler``,
+       ``process.refit_kalman``, ``model.cca``, ``model.refit_kalman``
+   * - ``sklearn``
+     - ``pandas``, ``river``, ``scikit-learn``
+     - ``process.adaptive_linear_regressor``, ``process.linear_regressor``,
+       ``process.sgd``, ``process.slda``, ``process.sklearn``, ``dim_reduce.*``
+   * - ``torch``
+     - ``torch``
+     - ``process.base``, ``process.torch``, ``process.rnn``,
+       ``process.transformer``, ``process.mlp_old``, ``model.mlp``,
+       ``model.rnn``, ``model.transformer``
+   * - ``all``
+     - both of the above
+     - everything, including all ``collection.sample_adapt_regressor`` backends
+
+:mod:`ezmsg.learn.collection.sample_adapt_regressor` imports its backend
+lazily, so it needs only the extra for the ``model_type`` in use — and none at
+all for ``model_type="kalman"``.
 
 Quick Start
 -----------

pyproject.toml

@@ -12,11 +12,26 @@ dependencies = [
     "ezmsg>=3.9.0",
     "ezmsg-baseproc>=1.7.0",
     "ezmsg-sigproc>=3.0.0",
+    "numpy",
+    "scipy",
+    "array-api-compat",
+]
+
+[project.optional-dependencies]
+# The heavyweight backends are opt-in so that the pure-numpy processors
+# (process.ssr, process.flatten, process.refit_kalman, ...) can be installed on
+# constrained hosts without pulling in a full PyTorch/sklearn stack.
+sklearn = [
     "pandas>=2.2",
     "river>=0.22.0",
     "scikit-learn>=1.6.0",
+]
+torch = [
     "torch>=2.6.0",
 ]
+all = [
+    "ezmsg-learn[sklearn,torch]",
+]
 
 [dependency-groups]
 dev = [
@@ -28,11 +43,17 @@ lint = [
     "ruff==0.16.2",
 ]
 test = [
+    # The backends are extras, so `uv sync` would not install them; the test
+    # suite covers all of them.
+    "ezmsg-learn[all]",
     "ezmsg-simbiophys>=1.8.0",
     "hmmlearn>=0.3.3",
     "pytest>=8.4.1",
 ]
 docs = [
+    # autodoc imports every documented module, so the docs build needs the
+    # project and all of its backends -- not just Sphinx.
+    "ezmsg-learn[all]",
     "sphinx>=8.1.3",
     "pydata-sphinx-theme",
     "sphinx_autodoc_typehints>=3.0.0",

src/ezmsg/learn/_optional.py

@@ -0,0 +1,25 @@
+"""Reporting for the optional backend dependencies.
+
+``ezmsg-learn`` ships its heavyweight backends as extras (``ezmsg-learn[torch]``
+and ``ezmsg-learn[sklearn]``) so the pure-numpy processors can be installed
+without them. Modules belonging to an extra still import their backend eagerly;
+they only wrap the import so that a missing install reports which extra to
+install instead of a bare ``ModuleNotFoundError: No module named 'torch'``.
+"""
+
+
+def missing_extra(extra: str, module: str) -> ImportError:
+    """Build the error raised when a module's optional backend is not installed.
+
+    Args:
+        extra: Name of the extra that provides the backend, e.g. ``"torch"``.
+        module: Importing module, normally passed as ``__name__``.
+
+    Returns:
+        The :class:`ImportError` to raise; chain it off the original with
+        ``raise missing_extra(...) from exc``.
+    """
+    return ImportError(
+        f"{module} requires the optional '{extra}' dependencies of ezmsg-learn, which are not installed. "
+        f'Install them with: pip install "ezmsg-learn[{extra}]"'
+    )

src/ezmsg/learn/collection/sample_adapt_regressor.py

@@ -1,3 +1,13 @@
+"""Decode collection wired around a single regressor engine.
+
+.. note::
+   The regressor backends live behind optional extras. ``model_type="mlp"``
+   needs ``ezmsg-learn[torch]`` and the River/sklearn model types need
+   ``ezmsg-learn[sklearn]``; ``model_type="kalman"`` needs neither. The backend
+   modules are imported only once a backend is selected, so an install carrying
+   just one extra can still build the collection for that backend.
+"""
+
 from dataclasses import field
 
 import ezmsg.core as ez
@@ -13,17 +23,12 @@
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ezmsg.learn.process.adaptive_linear_regressor import (
-    AdaptiveLinearRegressorSettings,
-    AdaptiveLinearRegressorUnit,
-)
 from ezmsg.learn.process.flatten import Flatten, FlattenSettings
 from ezmsg.learn.process.refit_kalman import (
     RefitKalmanFilterSettings,
     RefitKalmanFilterUnit,
 )
 from ezmsg.learn.process.seqseqsampler import SeqSeqSamplerSettings, SeqSeqSamplerUnit
-from ezmsg.learn.process.torch import TorchModelSettings, TorchModelUnit
 from ezmsg.learn.util import AdaptiveLinearRegressor
 
 #: Default torch model class used when ``model_type == "mlp"``.
@@ -181,16 +186,47 @@ def _build_regressor_unit(settings: SampleAdaptRegressorSettings):
     """Factory: construct the single regressor unit for ``settings.model_type``.
 
     Returns ``(unit, backend)`` where ``backend`` is ``"linear"`` (River/sklearn
-    via :class:`AdaptiveLinearRegressorUnit`), ``"torch"`` (mlp), or ``"kalman"``.
+    via ``AdaptiveLinearRegressorUnit``), ``"torch"`` (mlp), or ``"kalman"``.
     """
     backend = _model_backend(settings.model_type)
     if backend == "torch":
+        from ezmsg.learn.process.torch import TorchModelUnit
+
         return TorchModelUnit(), backend
     if backend == "kalman":
         return RefitKalmanFilterUnit(), backend
+
+    from ezmsg.learn.process.adaptive_linear_regressor import AdaptiveLinearRegressorUnit
+
     return AdaptiveLinearRegressorUnit(), backend
 
 
+def _build_regressor_settings(backend: str, settings: SampleAdaptRegressorSettings):
+    """Translate the collection settings into the selected backend's settings."""
+    if backend == "torch":
+        from ezmsg.learn.process.torch import TorchModelSettings
+
+        return TorchModelSettings(
+            model_class=settings.model_class,
+            checkpoint_path=settings.model_path,
+            model_kwargs=dict(settings.model_kwargs),
+            device=settings.device,
+        )
+    if backend == "kalman":
+        return RefitKalmanFilterSettings(
+            checkpoint_path=settings.model_path,
+            steady_state=settings.steady_state,
+        )
+
+    from ezmsg.learn.process.adaptive_linear_regressor import AdaptiveLinearRegressorSettings
+
+    return AdaptiveLinearRegressorSettings(
+        model_type=settings.model_type,
+        settings_path=settings.model_path,
+        model_kwargs=settings.model_kwargs,
+    )
+
+
 def build_sample_adapt_regressor(
     settings: SampleAdaptRegressorSettings,
 ) -> ez.Collection:
@@ -227,30 +263,7 @@ class SampleAdaptRegressor(ez.Collection):
             ADAPTER = DecodeOutputAdapter()
 
         def configure(self) -> None:
-            if backend == "linear":
-                self.REGRESSOR.apply_settings(
-                    AdaptiveLinearRegressorSettings(
-                        model_type=self.SETTINGS.model_type,
-                        settings_path=self.SETTINGS.model_path,
-                        model_kwargs=self.SETTINGS.model_kwargs,
-                    )
-                )
-            elif backend == "torch":
-                self.REGRESSOR.apply_settings(
-                    TorchModelSettings(
-                        model_class=self.SETTINGS.model_class,
-                        checkpoint_path=self.SETTINGS.model_path,
-                        model_kwargs=dict(self.SETTINGS.model_kwargs),
-                        device=self.SETTINGS.device,
-                    )
-                )
-            else:
-                self.REGRESSOR.apply_settings(
-                    RefitKalmanFilterSettings(
-                        checkpoint_path=self.SETTINGS.model_path,
-                        steady_state=self.SETTINGS.steady_state,
-                    )
-                )
+            self.REGRESSOR.apply_settings(_build_regressor_settings(backend, self.SETTINGS))
 
             if use_window:
                 self.WINDOW.apply_settings(

src/ezmsg/learn/dim_reduce/adaptive_decomp.py

@@ -19,7 +19,13 @@
     processor_state,
 )
 from ezmsg.util.messages.axisarray import AxisArray, replace
-from sklearn.decomposition import IncrementalPCA, MiniBatchNMF
+
+from .._optional import missing_extra
+
+try:
+    from sklearn.decomposition import IncrementalPCA, MiniBatchNMF
+except ImportError as exc:
+    raise missing_extra("sklearn", __name__) from exc
 
 
 class AdaptiveDecompSettings(ez.Settings):

src/ezmsg/learn/model/mlp.py

@@ -1,5 +1,10 @@
-import torch
-import torch.nn
+from .._optional import missing_extra
+
+try:
+    import torch
+    import torch.nn
+except ImportError as exc:
+    raise missing_extra("torch", __name__) from exc
 
 
 class MLP(torch.nn.Module):