Merge pull request #9 from adonath/safetensors_io

Add safetensors io

Author: adonath

gmmx/fit.py

@@ -161,4 +161,4 @@ def em_cond(
             body_fun=em_step,
             init_val=(x, gmm, 0, jnp.asarray(jnp.inf), jnp.array(jnp.inf)),
         )
-        return EMFitterResult(*result, converged=result[2] < self.max_iter)  # type: ignore [misc]
+        return EMFitterResult(*result, converged=result[2] < self.max_iter)

gmmx/gmm.py

@@ -48,6 +48,7 @@
 
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass, field
 from enum import Enum
 from functools import partial
@@ -69,6 +70,7 @@
     "GaussianMixtureSKLearn",
 ]
 
+log = logging.getLogger()
 
 AnyArray = Union[np.typing.NDArray, jax.Array]
 Device = Union[str, None]
@@ -663,6 +665,49 @@ def n_parameters(self) -> int:
             - 1
         )
 
+    def write(self, filename: str) -> None:
+        """Save the model parameters to a file in safetensors format."""
+        from safetensors.flax import save_file
+
+        data = {
+            "means": self.means_numpy,
+            "weights": self.weights_numpy,
+            "covariances": self.covariances.values_numpy,
+        }
+
+        metadata = {"covariance-type": self.covariances.type}
+
+        log.info(f"Writing {filename}")
+        save_file(data, metadata=metadata, filename=filename)  # type: ignore [arg-type]
+
+    @classmethod
+    def read(cls, filename: str, device: str = "cpu") -> GaussianMixtureModelJax:
+        """Read model parameters from a safetensors file.
+
+        Parameters
+        ----------
+        filename : str
+            Path to the safetensors file.
+        device : str, optional
+            Device to load the tensors onto (default: "cpu").
+
+        Returns
+        -------
+        GaussianMixtureModelJax
+            Loaded Gaussian Mixture Model instance.
+        """
+        from safetensors import safe_open
+
+        data = {}
+
+        with safe_open(filename, framework="flax", device=device) as f:
+            for key in f.keys():  # noqa: SIM118
+                data[key] = f.get_tensor(key)
+
+            covariance_type = f.metadata()["covariance-type"]
+
+        return cls.from_squeezed(**data, covariance_type=covariance_type)
+
     @property
     def log_weights(self) -> jax.Array:
         """Log weights (~jax.ndarray)"""

pyproject.toml

@@ -20,6 +20,7 @@ classifiers = [
 dependencies = [
     "jax>=0.4.30",
     "numpy>=1.26.0",
+    "safetensors>=0.5.0",
 ]
 
 [project.urls]
@@ -39,6 +40,7 @@ dev-dependencies = [
     "mkdocs-material>=8.5.10",
     "mkdocstrings[python]>=0.26.1",
     "scikit-learn>=1.0",
+    "safetensors>=0.5.0",
 ]
 
 [build-system]

tests/test_gmm.py

@@ -149,6 +149,22 @@ def test_fit(gmm_jax):
     assert_allclose(result.gmm.weights_numpy, [0.2, 0.8], rtol=0.05)
 
 
+def test_io(gmm_jax, tmpdir):
+    filename = tmpdir / "model.safetensors"
+
+    gmm_jax.write(filename)
+
+    new_model = GaussianMixtureModelJax.read(filename)
+
+    assert_allclose(gmm_jax.means_numpy, new_model.means_numpy)
+    assert_allclose(gmm_jax.weights_numpy, new_model.weights_numpy)
+    assert_allclose(
+        gmm_jax.covariances.values_numpy, new_model.covariances.values_numpy
+    )
+
+    assert gmm_jax.covariances.type == new_model.covariances.type
+
+
 def test_fit_against_sklearn(gmm_jax):
     # Fitting is hard to test, especillay we cannot guarantee the fit converges to the same solution
     # However the "global" likelihood (summed accross all components) for a given feature vector