Fix types

Author: pchlenski

manify/clustering/fuzzy_kmeans.py

@@ -120,7 +120,8 @@ def _init_centers(self, X: Float[torch.Tensor, "n_points n_features"]) -> None:
         # IMPORTANT: Use self.manifold.manifold for ManifoldParameter,
         # as self.manifold is our wrapper and self.manifold.manifold is the geoopt object.
         self.mu_ = ManifoldParameter(
-            centers.clone().detach(), manifold=self.manifold.manifold
+            centers.clone().detach(),  # type: ignore
+            manifold=self.manifold.manifold,
         )  # Ensure centers are detached
         self.mu_.requires_grad_(True)
 

manify/manifolds.py

@@ -829,7 +829,7 @@ def gaussian_mixture(
 
         # Generate cluster means
         cluster_means = self.sample(num_clusters, sigma_factorized=[torch.eye(M.dim) * cov_scale_means for M in self.P])
-        assert cluster_means.shape == (num_clusters, self.ambient_dim), "Cluster means shape mismatch."
+        assert cluster_means.shape == (num_clusters, self.ambient_dim), "Cluster means shape mismatch."  # type: ignore
 
         # Generate class assignments
         cluster_probs = torch.rand(num_clusters)