dataeval 0.72.1__py3-none-any.whl → 0.73.0__py3-none-any.whl

dataeval/{_internal/models/tensorflow → utils/tensorflow/_internal}/gmm.py

@@ -8,10 +8,16 @@ Licensed under Apache Software License (Apache 2.0)
 
 from __future__ import annotations
 
-from typing import NamedTuple
+from typing import TYPE_CHECKING, NamedTuple
 
 import numpy as np
-import tensorflow as tf
+
+from dataeval.utils.lazy import lazyload
+
+if TYPE_CHECKING:
+    import tensorflow as tf
+else:
+    tf = lazyload("tensorflow")
 
 
 class GaussianMixtureModelParams(NamedTuple):

dataeval/{_internal/models/tensorflow/losses.py → utils/tensorflow/_internal/loss.py}

@@ -8,16 +8,26 @@ Licensed under Apache Software License (Apache 2.0)
 
 from __future__ import annotations
 
-from typing import Literal, cast
+from typing import TYPE_CHECKING, Literal, cast
 
-import tensorflow as tf
+import numpy as np
 from numpy.typing import NDArray
-from tensorflow_probability.python.distributions.mvn_diag import MultivariateNormalDiag
-from tensorflow_probability.python.distributions.mvn_tril import MultivariateNormalTriL
-from tensorflow_probability.python.stats import covariance
-from tf_keras.layers import Flatten
 
-from dataeval._internal.models.tensorflow.gmm import gmm_energy, gmm_params
+from dataeval.utils.lazy import lazyload
+from dataeval.utils.tensorflow._internal.gmm import gmm_energy, gmm_params
+
+if TYPE_CHECKING:
+    import tensorflow as tf
+    import tensorflow_probability.python.distributions.mvn_diag as mvn_diag
+    import tensorflow_probability.python.distributions.mvn_tril as mvn_tril
+    import tensorflow_probability.python.stats as tfp_stats
+    import tf_keras as keras
+else:
+    tf = lazyload("tensorflow")
+    keras = lazyload("tf_keras")
+    mvn_diag = lazyload("tensorflow_probability.python.distributions.mvn_diag")
+    mvn_tril = lazyload("tensorflow_probability.python.distributions.mvn_tril")
+    tfp_stats = lazyload("tensorflow_probability.python.stats")
 
 
 class Elbo:
@@ -39,29 +49,29 @@ class Elbo:
     def __init__(
         self,
         cov_type: Literal["cov_full", "cov_diag"] | float = 1.0,
-        x: tf.Tensor | NDArray | None = None,
+        x: tf.Tensor | NDArray[np.float32] | None = None,
     ):
         if isinstance(cov_type, float):
-            self.cov = ("sim", cov_type)
+            self._cov = ("sim", cov_type)
         elif cov_type in ["cov_full", "cov_diag"]:
-            x_np: NDArray = x.numpy() if tf.is_tensor(x) else x  # type: ignore
-            cov = covariance(x_np.reshape(x_np.shape[0], -1))  # type: ignore py38
+            x_np: NDArray[np.float32] = x.numpy().astype(np.float32) if tf.is_tensor(x) else x  # type: ignore
+            cov = tfp_stats.covariance(x_np.reshape(x_np.shape[0], -1))  # type: ignore py38
             if cov_type == "cov_diag":  # infer standard deviation from covariance matrix
                 cov = tf.math.sqrt(tf.linalg.diag_part(cov))
-            self.cov = (cov_type, cov)
+            self._cov = (cov_type, cov)
         else:
             raise ValueError("Only cov_full, cov_diag or sim value should be specified.")
 
     def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
-        y_pred_flat = cast(tf.Tensor, Flatten()(y_pred))
+        y_pred_flat = cast(tf.Tensor, keras.layers.Flatten()(y_pred))
 
-        if self.cov[0] == "cov_full":
-            y_mn = MultivariateNormalTriL(y_pred_flat, scale_tril=tf.linalg.cholesky(self.cov[1]))
+        if self._cov[0] == "cov_full":
+            y_mn = mvn_tril.MultivariateNormalTriL(y_pred_flat, scale_tril=tf.linalg.cholesky(self._cov[1]))
         else:  # cov_diag and sim
-            cov_diag = self.cov[1] if self.cov[0] == "cov_diag" else self.cov[1] * tf.ones(y_pred_flat.shape[-1])
-            y_mn = MultivariateNormalDiag(y_pred_flat, scale_diag=cov_diag)
+            cov_diag = self._cov[1] if self._cov[0] == "cov_diag" else self._cov[1] * tf.ones(y_pred_flat.shape[-1])
+            y_mn = mvn_diag.MultivariateNormalDiag(y_pred_flat, scale_diag=cov_diag)
 
-        loss = -tf.reduce_mean(y_mn.log_prob(Flatten()(y_true)))
+        loss = -tf.reduce_mean(y_mn.log_prob(keras.layers.Flatten()(y_true)))
         return loss