def get_kmeans_module(cls, module: _torch.nn.Module) -> _Optional[_Type["KMeansModule"]]:
        """
        Returns the :py:class:`KMeansModule` class which implements k-means
        for the given module.
        """
        for _, layer_cls in cls.REGISTRY.items():
            if layer_cls.is_supported_module(module):
                return layer_cls
        return None

        Returns the :py:class:`KMeansModule` class which implements k-means
        for the given module.
        
coremltools/optimize/torch/_utils/k_means.py