"""
DKM: Differentiable K-Means Clustering Layer for Neural Network Compression

Implementation of the ICLR 2022 paper by Cho et al.
(arXiv: 2108.12659)
"""

from dkm.dkm_layer import DKMLayer
from dkm.compressor import DKMCompressor, compress_model
from dkm.utils import compute_model_size, compute_compression_ratio

__version__ = "1.0.0"
__all__ = [
    "DKMLayer",
    "DKMCompressor",
    "compress_model",
    "compute_model_size",
    "compute_compression_ratio",
]