Add to_device() method

Author: adonath

gmmx/gmm.py

@@ -665,6 +665,17 @@ def n_parameters(self) -> int:
             - 1
         )
 
+    def to_device(self, device: Any) -> GaussianMixtureModelJax:
+        """Move model to device"""
+
+        def move_array_to_device(node):  # type: ignore [no-untyped-def]
+            if isinstance(node, jax.Array):
+                return jax.device_put(node, device=device)
+
+            return node
+
+        return jax.tree_util.tree_map(move_array_to_device, self)  # type: ignore [no-any-return]
+
     def write(self, filename: str) -> None:
         """Save the model parameters to a file in safetensors format."""
         from safetensors.flax import save_file

tests/test_gmm.py

@@ -149,6 +149,12 @@ def test_fit(gmm_jax):
     assert_allclose(result.gmm.weights_numpy, [0.2, 0.8], rtol=0.05)
 
 
+def test_move_to_device(gmm_jax):
+    gmm_jax_moved = gmm_jax.to_device(jax.devices()[0])
+    assert_allclose(gmm_jax.means_numpy, gmm_jax_moved.means_numpy)
+    assert gmm_jax_moved.means.devices() == {jax.devices()[0]}
+
+
 def test_io(gmm_jax, tmpdir):
     filename = tmpdir / "model.safetensors"