Fix type annotatoions

Author: adonath

gmmx/gmm.py

@@ -426,7 +426,7 @@ def log_prob(self, x: jax.Array, means: jax.Array) -> jax.Array:
         )
 
     @classmethod
-    def from_precisions(cls, precisions: jax.AnyArray) -> DiagCovariances:
+    def from_precisions(cls, precisions: AnyArray) -> DiagCovariances:
         """Create covariance matrix from precision matrices"""
         values = 1.0 / precisions
         return cls.from_squeezed(values=values)
@@ -964,7 +964,7 @@ def _initialize_gmm(self, x: AnyArray) -> None:
             self._gmm = GaussianMixtureModelJax.from_squeezed(
                 means=self.means_init,  # type: ignore [arg-type]
                 covariances=covar.from_precisions(
-                    self.precisions_init.astype(np.float32)
+                    self.precisions_init.astype(np.float32)  # type: ignore [union-attr]
                 ).values_numpy,
                 weights=self.weights_init,  # type: ignore [arg-type]
                 covariance_type=self.covariance_type,