torchgmm 0.1.0__tar.gz

torchgmm-0.1.0/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2020 Oliver Borchert
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchgmm-0.1.0/PKG-INFO

@@ -0,0 +1,138 @@
+Metadata-Version: 2.1
+Name: torchgmm
+Version: 0.1.0
+Summary: Gaussian Mixture Models in PyTorch.
+Home-page: https://github.com/marcovarrone/torchgmm
+License: MIT
+Author: Oliver Borchert
+Author-email: me@borchero.com
+Requires-Python: >=3.8,<4.0
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: numpy (>=1.20.3,<2.0.0)
+Requires-Dist: pytorch-lightning (>2.0)
+Requires-Dist: torch (>1.11.0)
+Requires-Dist: torchmetrics (>=0.6)
+Project-URL: Repository, https://github.com/marcovarrone/torchgmm
+Description-Content-Type: text/markdown
+
+# TorchGMM
+
+<!-- ![PyPi](https://img.shields.io/pypi/v/torchgmm?label=version)
+![License](https://img.shields.io/pypi/l/torchgmm) -->
+
+TorchGMM allows to run Gaussian Mixture Models on single or multiple CPUs/GPUs.
+The repository is a fork from [PyCave](https://github.com/borchero/pycave) and [LightKit](https://github.com/borchero/lightkit), two amazing packages developed by [Olivier Borchert](https://github.com/borchero) that are not being maintained anymore.
+While PyCave implements additional models such as Markov Chains, TorchGMM focuses only on Gaussian Mixture Models.
+
+The models are implemented in [PyTorch](https://pytorch.org/) and [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/), and provide an `Estimator` API
+that is fully compatible with [scikit-learn](https://scikit-learn.org/stable/).
+
+For Gaussian mixture model, TorchGMM allows for 100x speed ups when using a GPU and enables to train
+on markedly larger datasets via mini-batch training. The full suite of benchmarks run to compare
+TorchGMM models against scikit-learn models is available on the
+[documentation website](https://pycave.borchero.com/sites/benchmark.html).
+
+## Features
+
+- Support for GPU and multi-node training by implementing models in PyTorch and relying on
+  [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/)
+- Mini-batch training for all models such that they can be used on huge datasets
+- Well-structured implementation of models
+
+  - High-level `Estimator` API allows for easy usage such that models feel and behave like in
+    scikit-learn
+  - Medium-level `LightingModule` implements the training algorithm
+  - Low-level PyTorch `Module` manages the model parameters
+
+## Installation
+
+TorchGMM is available via `pip`:
+
+```bash
+pip install torchgmm
+```
+
+If you are using [Poetry](https://python-poetry.org/):
+
+```bash
+poetry add torchgmm
+```
+
+## Usage
+
+If you've ever used scikit-learn, you'll feel right at home when using TorchGMM. First, let's create
+some artificial data to work with:
+
+```python
+import torch
+
+X = torch.cat([
+    torch.randn(10000, 8) - 5,
+    torch.randn(10000, 8),
+    torch.randn(10000, 8) + 5,
+])
+```
+
+This dataset consists of three clusters with 8-dimensional datapoints. If you want to fit a K-Means
+model, to find the clusters' centroids, it's as easy as:
+
+```python
+from torchgmm.clustering import KMeans
+
+estimator = KMeans(3)
+estimator.fit(X)
+
+# Once the estimator is fitted, it provides various properties. One of them is
+# the `model_` property which yields the PyTorch module with the fitted parameters.
+print("Centroids are:")
+print(estimator.model_.centroids)
+```
+
+Due to the high-level estimator API, the usage for all machine learning models is similar. The API
+documentation provides more detailed information about parameters that can be passed to estimators
+and which methods are available.
+
+### GPU and Multi-Node training
+
+For GPU- and multi-node training, TorchGMM leverages PyTorch Lightning. The hardware that training
+runs on is determined by the
+[Trainer](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.html#pytorch_lightning.trainer.trainer.Trainer)
+class. It's
+[**init**](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.html#pytorch_lightning.trainer.trainer.Trainer.__init__)
+method provides various configuration options.
+
+If you want to run K-Means with a GPU, you can pass the options `accelerator='gpu'` and `devices=1`
+to the estimator's initializer:
+
+```python
+estimator = KMeans(3, trainer_params=dict(accelerator='gpu', devices=1))
+```
+
+Similarly, if you want to train on 4 nodes simultaneously where each node has one GPU available,
+you can specify this as follows:
+
+```python
+estimator = KMeans(3, trainer_params=dict(num_nodes=4, accelerator='gpu', devices=1))
+```
+
+In fact, **you do not need to change anything else in your code**.
+
+### Implemented Models
+
+Currently, TorchGMM implements two different models:
+
+- [GaussianMixture](https://pycave.borchero.com/sites/generated/bayes/gmm/pycave.bayes.GaussianMixture.html)
+- [K-Means](https://pycave.borchero.com/sites/generated/clustering/kmeans/pycave.clustering.KMeans.html)
+
+## License
+
+TorchGMM is licensed under the [MIT License](https://github.com/marcovarrone/torchgmm/blob/main/LICENSE).
+

torchgmm-0.1.0/README.md

@@ -0,0 +1,112 @@
+# TorchGMM
+
+<!-- ![PyPi](https://img.shields.io/pypi/v/torchgmm?label=version)
+![License](https://img.shields.io/pypi/l/torchgmm) -->
+
+TorchGMM allows to run Gaussian Mixture Models on single or multiple CPUs/GPUs.
+The repository is a fork from [PyCave](https://github.com/borchero/pycave) and [LightKit](https://github.com/borchero/lightkit), two amazing packages developed by [Olivier Borchert](https://github.com/borchero) that are not being maintained anymore.
+While PyCave implements additional models such as Markov Chains, TorchGMM focuses only on Gaussian Mixture Models.
+
+The models are implemented in [PyTorch](https://pytorch.org/) and [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/), and provide an `Estimator` API
+that is fully compatible with [scikit-learn](https://scikit-learn.org/stable/).
+
+For Gaussian mixture model, TorchGMM allows for 100x speed ups when using a GPU and enables to train
+on markedly larger datasets via mini-batch training. The full suite of benchmarks run to compare
+TorchGMM models against scikit-learn models is available on the
+[documentation website](https://pycave.borchero.com/sites/benchmark.html).
+
+## Features
+
+- Support for GPU and multi-node training by implementing models in PyTorch and relying on
+  [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/)
+- Mini-batch training for all models such that they can be used on huge datasets
+- Well-structured implementation of models
+
+  - High-level `Estimator` API allows for easy usage such that models feel and behave like in
+    scikit-learn
+  - Medium-level `LightingModule` implements the training algorithm
+  - Low-level PyTorch `Module` manages the model parameters
+
+## Installation
+
+TorchGMM is available via `pip`:
+
+```bash
+pip install torchgmm
+```
+
+If you are using [Poetry](https://python-poetry.org/):
+
+```bash
+poetry add torchgmm
+```
+
+## Usage
+
+If you've ever used scikit-learn, you'll feel right at home when using TorchGMM. First, let's create
+some artificial data to work with:
+
+```python
+import torch
+
+X = torch.cat([
+    torch.randn(10000, 8) - 5,
+    torch.randn(10000, 8),
+    torch.randn(10000, 8) + 5,
+])
+```
+
+This dataset consists of three clusters with 8-dimensional datapoints. If you want to fit a K-Means
+model, to find the clusters' centroids, it's as easy as:
+
+```python
+from torchgmm.clustering import KMeans
+
+estimator = KMeans(3)
+estimator.fit(X)
+
+# Once the estimator is fitted, it provides various properties. One of them is
+# the `model_` property which yields the PyTorch module with the fitted parameters.
+print("Centroids are:")
+print(estimator.model_.centroids)
+```
+
+Due to the high-level estimator API, the usage for all machine learning models is similar. The API
+documentation provides more detailed information about parameters that can be passed to estimators
+and which methods are available.
+
+### GPU and Multi-Node training
+
+For GPU- and multi-node training, TorchGMM leverages PyTorch Lightning. The hardware that training
+runs on is determined by the
+[Trainer](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.html#pytorch_lightning.trainer.trainer.Trainer)
+class. It's
+[**init**](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.html#pytorch_lightning.trainer.trainer.Trainer.__init__)
+method provides various configuration options.
+
+If you want to run K-Means with a GPU, you can pass the options `accelerator='gpu'` and `devices=1`
+to the estimator's initializer:
+
+```python
+estimator = KMeans(3, trainer_params=dict(accelerator='gpu', devices=1))
+```
+
+Similarly, if you want to train on 4 nodes simultaneously where each node has one GPU available,
+you can specify this as follows:
+
+```python
+estimator = KMeans(3, trainer_params=dict(num_nodes=4, accelerator='gpu', devices=1))
+```
+
+In fact, **you do not need to change anything else in your code**.
+
+### Implemented Models
+
+Currently, TorchGMM implements two different models:
+
+- [GaussianMixture](https://pycave.borchero.com/sites/generated/bayes/gmm/pycave.bayes.GaussianMixture.html)
+- [K-Means](https://pycave.borchero.com/sites/generated/clustering/kmeans/pycave.clustering.KMeans.html)
+
+## License
+
+TorchGMM is licensed under the [MIT License](https://github.com/marcovarrone/torchgmm/blob/main/LICENSE).

torchgmm-0.1.0/pyproject.toml

@@ -0,0 +1,95 @@
+[tool.poetry]
+authors = ["Oliver Borchert <me@borchero.com>"]
+classifiers = [
+  "Development Status :: 5 - Production/Stable",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+description = "Gaussian Mixture Models in PyTorch."
+license = "MIT"
+name = "torchgmm"
+readme = "README.md"
+repository = "https://github.com/marcovarrone/torchgmm"
+version = "0.1.0"
+
+[tool.poetry.dependencies]
+numpy = "^1.20.3"
+python = ">=3.8,<4.0"
+pytorch-lightning = ">2.0"
+torch = ">1.11.0"
+torchmetrics = ">=0.6,"
+
+[tool.poetry.group.pre-commit.dependencies]
+black = "^22.12.0"
+docformatter = "^1.5.0"
+isort = "^5.10.1"
+mypy = "^0.991"
+pylint = "^2.12.2"
+pyupgrade = "^3.3.1"
+
+[tool.poetry.group.docs.dependencies]
+Sphinx = "^5.0.0"
+pydata-sphinx-theme = "^0.7.2"
+scanpydoc = "^0.7.1"
+sphinx-autodoc-typehints = "^1.12.0"
+sphinx-automodapi = "^0.13"
+sphinx-copybutton = "^0.3.3"
+sphinxcontrib-spelling = "^7.2.1"
+
+[tool.poetry.group.testing.dependencies]
+flaky = "^3.7.0"
+pytest = "^6.2.4"
+pytest-benchmark = "^3.4.1"
+scikit-learn = ">=0.24.2"
+
+[tool.poetry.group.dev.dependencies]
+jupyter = "^1.0.0"
+
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core>=1.0.0"]
+
+[tool.pylint.messages_control]
+disable = [
+  "arguments-differ",
+  "duplicate-code",
+  "missing-module-docstring",
+  "invalid-name",
+  "too-few-public-methods",
+  "too-many-ancestors",
+  "too-many-arguments",
+  "too-many-branches",
+  "too-many-locals",
+  "too-many-instance-attributes",
+]
+
+[tool.pylint.typecheck]
+generated-members = [
+  "torch.*",
+]
+
+[tool.black]
+line-length = 99
+target-version = ["py38", "py39", "py310"]
+
+[tool.isort]
+force_alphabetical_sort_within_sections = true
+include_trailing_comma = true
+known_first_party = "torchgmm,tests"
+line_length = 99
+lines_between_sections = 0
+profile = "black"
+skip_gitignore = true
+
+[tool.docformatter]
+make-summary-multi-line = true
+pre-summary-newline = true
+recursive = true
+wrap-descriptions = 99
+wrap-summaries = 99
+
+[tool.pytest.ini_options]
+filterwarnings = [
+  "ignore:.*Create unlinked descriptors is going to go away.*:DeprecationWarning",
+  "ignore:.*this fit will run with no optimizer.*",
+  "ignore:.*Consider increasing the value of the `num_workers` argument.*",
+]

torchgmm-0.1.0/torchgmm/__init__.py

@@ -0,0 +1,38 @@
+import logging
+import warnings
+import torchgmm.base
+
+# This is taken from PyTorch Lightning and ensures that logging for this package is enabled
+_root_logger = logging.getLogger()
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+if not _root_logger.hasHandlers():
+    _logger.addHandler(logging.StreamHandler())
+    _logger.propagate = False
+
+# This disables most logs generated by PyTorch Lightning
+logging.getLogger("pytorch_lightning").setLevel(logging.WARNING)
+warnings.filterwarnings(
+    action="ignore", message=".*Consider increasing the value of the `num_workers` argument.*"
+)
+warnings.filterwarnings(
+    action="ignore", message=".*`LightningModule.configure_optimizers` returned `None`.*"
+)
+warnings.filterwarnings(
+    action="ignore", message=".*`LoggerConnector.gpus_metrics` was deprecated in v1.5.*"
+)
+
+# We also want to define a function which silences info logs
+def set_logging_level(level: int) -> None:
+    """
+    Enables or disables logging for the entire module. By default, logging is enabled.
+
+    Args:
+        enabled: Whether to enable logging.
+    """
+    _logger.setLevel(level)
+    torchgmm.base.set_logging_level(level)
+
+
+# Export
+__all__ = ["set_logging_level"]

torchgmm-0.1.0/torchgmm/base/__init__.py

@@ -0,0 +1,23 @@
+import logging
+from .estimator import BaseEstimator, ConfigurableBaseEstimator
+
+# This is taken from PyTorch Lightning and ensures that logging for this package is enabled
+_root_logger = logging.getLogger()
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+if not _root_logger.hasHandlers():
+    _logger.addHandler(logging.StreamHandler())
+    _logger.propagate = False
+
+
+def set_logging_level(level: int) -> None:
+    """
+    Enables or disables logging for the entire module. By default, logging is enabled.
+
+    Args:
+        enabled: Whether to enable logging.
+    """
+    _logger.setLevel(level)
+
+
+__all__ = ["BaseEstimator", "ConfigurableBaseEstimator"]

torchgmm-0.1.0/torchgmm/base/data/__init__.py

@@ -0,0 +1,14 @@
+from .collation import collate_tensor, collate_tuple
+from .loader import DataLoader
+from .sampler import RangeBatchSampler
+from .types import DataLoaderLike, dataset_from_tensors, TensorLike
+
+__all__ = [
+    "DataLoader",
+    "DataLoaderLike",
+    "RangeBatchSampler",
+    "TensorLike",
+    "collate_tensor",
+    "collate_tuple",
+    "dataset_from_tensors",
+]

torchgmm-0.1.0/torchgmm/base/data/collation.py

@@ -0,0 +1,22 @@
+from typing import Tuple
+import torch
+
+
+def collate_tuple(batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
+    """
+    Collate a tuple of batch items by returning the input tuple.
+
+    This is the default used by
+    :class:`~torchgmm.base.data.DataLoader` when slices are cut from the underlying data source.
+    """
+    return batch
+
+
+def collate_tensor(batch: Tuple[torch.Tensor, ...]) -> torch.Tensor:
+    """
+    Collates a tuple of batch items into the first tensor.
+
+    Might be useful if only a single tensor is passed to
+    :class:`~torchgmm.base.data.DataLoader`.
+    """
+    return batch[0]

torchgmm-0.1.0/torchgmm/base/data/loader.py

@@ -0,0 +1,78 @@
+from typing import Any, Iterator, TypeVar
+from pytorch_lightning.overrides.distributed import _IndexBatchSamplerWrapper
+from torch.utils.data import DataLoader as TorchDataLoader
+from torch.utils.data import Dataset, TensorDataset
+from torch.utils.data.sampler import SequentialSampler
+from .collation import collate_tuple
+from .sampler import RangeBatchSampler
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class DataLoader(TorchDataLoader[T_co]):
+    """Extension for PyTorch's builtin dataloader. This implementation allows
+    to retrieve contiguous indices from a
+    :class:`~torch.utils.data.TensorDataset` orders of magnitude faster. The
+    data loader, thus, enables to implement traditional machine learning
+    methods that exhibit a speed similar to the implementations found in
+    Scikit-learn.
+
+    Note:
+        Retrieving contiguous indices is only possible when all of the following conditions apply:
+
+        - ``shuffle=False`` or ``batch_sampler`` is of type
+          :class:`~torchgmm.base.data.RangeBatchSampler`
+        - ``sampler is None``
+        - ``num_workers=0``
+        - ``dataset`` is not iterable
+    """
+
+    def __init__(self, dataset: Dataset[T_co], **kwargs: Any):
+        """
+        Args:
+            dataset: The dataset from which to load the data.
+            kwargs: Keyword arguments passed to :meth:`torch.utils.data.DataLoader.__init__`.
+        """
+        if (
+            not kwargs.get("shuffle", False)
+            and "sampler" not in kwargs
+            and "batch_sampler" not in kwargs
+            and kwargs.get("num_workers", 0) == 0
+            and isinstance(dataset, TensorDataset)
+        ):
+            kwargs["batch_sampler"] = RangeBatchSampler(
+                SequentialSampler(dataset),
+                batch_size=kwargs.get("batch_size", 1),
+                drop_last=kwargs.get("drop_last", False),
+            )
+            kwargs.pop("batch_size", None)
+            kwargs.pop("shuffle", None)
+            kwargs.pop("drop_last", None)
+            kwargs.setdefault("collate_fn", collate_tuple)
+
+        super().__init__(dataset, **kwargs)  # type: ignore
+
+        self.custom_batching = self.num_workers == 0 and (
+            isinstance(self.batch_sampler, RangeBatchSampler)
+            or (
+                self.batch_sampler is not None
+                and hasattr(self.batch_sampler, "sampler")
+                and isinstance(self.batch_sampler.sampler, RangeBatchSampler)
+            )
+            or (
+                isinstance(self.batch_sampler, _IndexBatchSamplerWrapper)
+                and isinstance(self.batch_sampler._batch_sampler, RangeBatchSampler)  # type: ignore
+            )
+        )
+
+    def __iter__(self) -> Iterator[Any]:  # pylint: disable=inconsistent-return-statements
+        if not self.custom_batching:
+            yield from super().__iter__()
+            return
+
+        for indices in self.batch_sampler:
+            if isinstance(indices, range):
+                subscript = slice(indices.start, indices.stop)
+                yield self.collate_fn(tuple(t[subscript] for t in self.dataset.tensors))
+            else:
+                yield self.collate_fn(tuple(t[indices] for t in self.dataset.tensors))

torchgmm-0.1.0/torchgmm/base/data/sampler.py

@@ -0,0 +1,41 @@
+import math
+from typing import Iterator
+from torch.utils.data import Sampler
+from torch.utils.data.sampler import SequentialSampler
+
+
+class RangeBatchSampler(Sampler[range]):
+    """
+    Sampler providing batches of contiguous indices.
+
+    This sampler can be used with
+    :class:`torchgmm.base.data.DataLoader` to provide significant speedups for tensor datasets.
+    """
+
+    def __init__(self, sampler: Sampler[int], batch_size: int, drop_last: bool = False):
+        """
+        Args:
+            sampler: The sampler providing indices. Must be a sequential sampler. Note that the
+                only purpose of this sampler is to determine its length.
+            batch_size: The number of items to sample for each batch.
+            drop_last: Whether to drop the last batch if ``num_items`` is not divisible by
+                ``batch_size``.
+        """
+        assert isinstance(
+            sampler, SequentialSampler
+        ), f"{self.__class__.__name__} only works with sequential samplers."
+
+        super().__init__(None)
+        self.dataset_size = len(sampler)
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+
+    def __len__(self) -> int:
+        if self.drop_last:
+            return self.dataset_size // self.batch_size
+        return math.ceil(self.dataset_size / self.batch_size)
+
+    def __iter__(self) -> Iterator[range]:
+        for i in range(len(self)):
+            sample = range(i * self.batch_size, (i + 1) * self.batch_size)
+            yield sample

torchgmm-0.1.0/torchgmm/base/data/types.py

@@ -0,0 +1,46 @@
+from typing import Sequence, TypeVar, Union
+import numpy as np
+import numpy.typing as npt
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+T = TypeVar("T")
+
+TensorLike = Union[
+    npt.NDArray[np.float32],
+    torch.Tensor,
+]
+TensorLike.__doc__ = """
+Type annotation for functions accepting any kind of tensor data as input. Consider using this
+annotation if your methods in an estimator derived from :class:`torchgmm.base.BaseEstimator` work on
+tensors.
+"""
+
+DataLoaderLike = Union[
+    DataLoader[T],
+    Sequence[DataLoader[T]],
+]
+DataLoaderLike.__doc__ = """
+Generic type annotation for functions accepting any data loader as input. Consider using this
+annotation for the implementation of methods in an estimator derived from
+:class:`torchgmm.base.BaseEstimator`.
+"""
+
+
+def dataset_from_tensors(*data: TensorLike) -> TensorDataset:
+    """
+    Transforms a set of tensor-like items into a datasets.
+
+    Args:
+        data: The tensor-like items.
+
+    Returns:
+        The dataset.
+    """
+    return TensorDataset(*[_to_tensor(t) for t in data])
+
+
+def _to_tensor(data: TensorLike) -> torch.Tensor:
+    if isinstance(data, np.ndarray):
+        return torch.from_numpy(data)
+    return data

torchgmm-0.1.0/torchgmm/base/estimator/__init__.py

@@ -0,0 +1,12 @@
+from .base import BaseEstimator
+from .configurable import ConfigurableBaseEstimator
+from .exception import NotFittedError
+from .mixins import PredictorMixin, TransformerMixin
+
+__all__ = [
+    "BaseEstimator",
+    "ConfigurableBaseEstimator",
+    "NotFittedError",
+    "PredictorMixin",
+    "TransformerMixin",
+]

torchgmm-0.1.0/torchgmm/base/estimator/_protocols.py

@@ -0,0 +1,22 @@
+# pylint: disable=missing-class-docstring,missing-function-docstring
+
+from typing import Protocol, TypeVar
+
+D_contra = TypeVar("D_contra", contravariant=True)
+R_co = TypeVar("R_co", covariant=True)
+E = TypeVar("E", bound="Estimator")  # type: ignore
+
+
+class Estimator(Protocol[D_contra]):
+    def fit(self: E, data: D_contra) -> E:
+        ...
+
+
+class Transformer(Estimator[D_contra], Protocol[D_contra, R_co]):
+    def transform(self, data: D_contra) -> R_co:
+        ...
+
+
+class Predictor(Estimator[D_contra], Protocol[D_contra, R_co]):
+    def predict(self, data: D_contra) -> R_co:
+        ...