dataeval 0.73.1__py3-none-any.whl → 0.74.0__py3-none-any.whl

dataeval/metrics/bias/metadata_utils.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+__all__ = []
+
+import contextlib
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+from dataeval.interop import to_numpy
+
+with contextlib.suppress(ImportError):
+    from matplotlib.figure import Figure
+
+
+def get_counts(data: NDArray[np.int_], min_num_bins: int | None = None) -> NDArray[np.int_]:
+    """
+    Returns columnwise unique counts for discrete data.
+
+    Parameters
+    ----------
+    data : NDArray
+        Array containing integer values for metadata factors
+    min_num_bins : int | None, default None
+        Minimum number of bins for bincount, helps force consistency across runs
+
+    Returns
+    -------
+    NDArray[np.int_]
+        Bin counts per column of data.
+    """
+    max_value = data.max() + 1 if min_num_bins is None else min_num_bins
+    cnt_array = np.zeros((max_value, data.shape[1]), dtype=np.int_)
+    for idx in range(data.shape[1]):
+        cnt_array[:, idx] = np.bincount(data[:, idx], minlength=max_value)
+
+    return cnt_array
+
+
+def heatmap(
+    data: ArrayLike,
+    row_labels: list[str] | ArrayLike,
+    col_labels: list[str] | ArrayLike,
+    xlabel: str = "",
+    ylabel: str = "",
+    cbarlabel: str = "",
+) -> Figure:
+    """
+    Plots a formatted heatmap
+
+    Parameters
+    ----------
+    data : NDArray
+        Array containing numerical values for factors to plot
+    row_labels : ArrayLike
+        List/Array containing the labels for rows in the histogram
+    col_labels : ArrayLike
+        List/Array containing the labels for columns in the histogram
+    xlabel : str, default ""
+        X-axis label
+    ylabel : str, default ""
+        Y-axis label
+    cbarlabel : str, default ""
+        Label for the colorbar
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        Formatted heatmap
+    """
+    import matplotlib.pyplot as plt
+    from matplotlib.ticker import FuncFormatter
+
+    np_data = to_numpy(data)
+    rows = row_labels if isinstance(row_labels, list) else to_numpy(row_labels)
+    cols = col_labels if isinstance(col_labels, list) else to_numpy(col_labels)
+
+    fig, ax = plt.subplots(figsize=(10, 10))
+
+    # Plot the heatmap
+    im = ax.imshow(np_data, vmin=0, vmax=1.0)
+
+    # Create colorbar
+    cbar = fig.colorbar(im, shrink=0.5)
+    cbar.set_ticks([0.0, 0.25, 0.5, 0.75, 1.0])
+    cbar.set_ticklabels(["0.0", "0.25", "0.5", "0.75", "1.0"])
+    cbar.set_label(cbarlabel, loc="center")
+
+    # Show all ticks and label them with the respective list entries.
+    ax.set_xticks(np.arange(np_data.shape[1]), labels=cols)
+    ax.set_yticks(np.arange(np_data.shape[0]), labels=rows)
+
+    ax.tick_params(top=False, bottom=True, labeltop=False, labelbottom=True)
+    # Rotate the tick labels and set their alignment.
+    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
+
+    # Turn spines off and create white grid.
+    ax.spines[:].set_visible(False)
+
+    ax.set_xticks(np.arange(np_data.shape[1] + 1) - 0.5, minor=True)
+    ax.set_yticks(np.arange(np_data.shape[0] + 1) - 0.5, minor=True)
+    ax.grid(which="minor", color="w", linestyle="-", linewidth=3)
+    ax.tick_params(which="minor", bottom=False, left=False)
+
+    if xlabel:
+        ax.set_xlabel(xlabel)
+    if ylabel:
+        ax.set_ylabel(ylabel)
+
+    valfmt = FuncFormatter(format_text)
+
+    # Normalize the threshold to the images color range.
+    threshold = im.norm(1.0) / 2.0
+
+    # Set default alignment to center, but allow it to be
+    # overwritten by textkw.
+    kw = {"horizontalalignment": "center", "verticalalignment": "center"}
+
+    # Loop over the data and create a `Text` for each "pixel".
+    # Change the text's color depending on the data.
+    textcolors = ("white", "black")
+    texts = []
+    for i in range(np_data.shape[0]):
+        for j in range(np_data.shape[1]):
+            kw.update(color=textcolors[int(im.norm(np_data[i, j]) > threshold)])
+            text = im.axes.text(j, i, valfmt(np_data[i, j], None), **kw)  # type: ignore
+            texts.append(text)
+
+    fig.tight_layout()
+    return fig
+
+
+# Function to define how the text is displayed in the heatmap
+def format_text(*args: str) -> str:
+    """
+    Helper function to format text for heatmap()
+
+    Parameters
+    ----------
+    *args : tuple[str, str]
+        Text to be formatted. Second element is ignored, but is a
+        mandatory pass-through argument as per matplotlib.ticket.FuncFormatter
+
+    Returns
+    -------
+    str
+        Formatted text
+    """
+    x = args[0]
+    return f"{x:.2f}".replace("0.00", "0").replace("0.", ".").replace("nan", "")
+
+
+def diversity_bar_plot(labels: NDArray[Any], bar_heights: NDArray[Any]) -> Figure:
+    """
+    Plots a formatted bar plot
+
+    Parameters
+    ----------
+    labels : NDArray
+        Array containing the labels for each bar
+    bar_heights : NDArray
+        Array containing the values for each bar
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        Bar plot figure
+    """
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(figsize=(10, 10))
+
+    ax.bar(labels, bar_heights)
+    ax.set_xlabel("Factors")
+
+    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
+
+    fig.tight_layout()
+    return fig
+
+
+def coverage_plot(images: NDArray[Any], num_images: int) -> Figure:
+    """
+    Creates a single plot of all of the provided images
+
+    Parameters
+    ----------
+    images : NDArray
+        Array containing only the desired images to plot
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        Plot of all provided images
+    """
+    import matplotlib.pyplot as plt
+
+    num_images = min(num_images, len(images))
+
+    if images.ndim == 4:
+        images = np.moveaxis(images, 1, -1)
+    elif images.ndim == 3:
+        images = np.repeat(images[:, :, :, np.newaxis], 3, axis=-1)
+    else:
+        raise ValueError(
+            f"Expected a (N,C,H,W) or a (N, H, W) set of images, but got a {images.ndim}-dimensional set of images."
+        )
+
+    rows = int(np.ceil(num_images / 3))
+    fig, axs = plt.subplots(rows, 3, figsize=(9, 3 * rows))
+
+    if rows == 1:
+        for j in range(3):
+            if j >= len(images):
+                continue
+            axs[j].imshow(images[j])
+            axs[j].axis("off")
+    else:
+        for i in range(rows):
+            for j in range(3):
+                i_j = i * 3 + j
+                if i_j >= len(images):
+                    continue
+                axs[i, j].imshow(images[i_j])
+                axs[i, j].axis("off")
+
+    fig.tight_layout()
+    return fig

dataeval/metrics/bias/parity.py

@@ -4,14 +4,15 @@ __all__ = ["ParityOutput", "parity", "label_parity"]
 
 import warnings
 from dataclasses import dataclass
-from typing import Any, Generic, Mapping, TypeVar
+from typing import Any, Generic, TypeVar
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
-from scipy.stats import chi2_contingency, chisquare
+from scipy.stats import chisquare
+from scipy.stats.contingency import chi2_contingency, crosstab
 
-from dataeval.interop import to_numpy
-from dataeval.metrics.bias.metadata import CLASS_LABEL, preprocess_metadata
+from dataeval.interop import as_numpy, to_numpy
+from dataeval.metrics.bias.metadata_preprocessing import MetadataOutput
 from dataeval.output import OutputMetadata, set_metadata
 
 TData = TypeVar("TData", np.float64, NDArray[np.float64])
@@ -37,97 +38,6 @@ class ParityOutput(Generic[TData], OutputMetadata):
     metadata_names: list[str] | None
 
 
-def digitize_factor_bins(continuous_values: NDArray[Any], bins: int, factor_name: str) -> NDArray[np.intp]:
-    """
-    Digitizes a list of values into a given number of bins.
-
-    Parameters
-    ----------
-    continuous_values : NDArray
-        The values to be digitized.
-    bins : int
-        The number of bins for the discrete values that continuous_values will be digitized into.
-    factor_name : str
-        The name of the factor to be digitized.
-
-    Returns
-    -------
-    NDArray[np.intp]
-        The digitized values
-    """
-
-    if not np.all([np.issubdtype(type(n), np.number) for n in continuous_values]):
-        raise TypeError(
-            f"Encountered a non-numeric value for factor {factor_name}, but the factor"
-            " was specified to be continuous. Ensure all occurrences of this factor are numeric types,"
-            f" or do not specify {factor_name} as a continuous factor."
-        )
-
-    _, bin_edges = np.histogram(continuous_values, bins=bins)
-    bin_edges[-1] = np.inf
-    bin_edges[0] = -np.inf
-    return np.digitize(continuous_values, bin_edges)
-
-
-def format_discretize_factors(
-    data: NDArray[Any],
-    names: list[str],
-    is_categorical: list[bool],
-    continuous_factor_bincounts: Mapping[str, int] | None,
-) -> dict[str, NDArray[Any]]:
-    """
-    Sets up the internal list of metadata factors.
-
-    Parameters
-    ----------
-    data : NDArray
-        The dataset factors, which are per-image attributes including class label and metadata.
-    names : list[str]
-        The class label
-    continuous_factor_bincounts : Mapping[str, int] or None
-        The factors in data_factors that have continuous values and the array of bin counts to
-        discretize values into. All factors are treated as having discrete values unless they
-        are specified as keys in this dictionary. Each element of this array must occur as a key
-        in data_factors.
-
-    Returns
-    -------
-    Dict[str, NDArray]
-        - Intrinsic per-image metadata information with the formatting that input data_factors uses.
-          Each key is a metadata factor, whose value is the discrete per-image factor values.
-    """
-
-    if continuous_factor_bincounts:
-        invalid_keys = set(continuous_factor_bincounts.keys()) - set(names)
-        if invalid_keys:
-            raise KeyError(
-                f"The continuous factor(s) {invalid_keys} do not exist in data_factors. Delete these "
-                "keys from `continuous_factor_names` or add corresponding entries to `data_factors`."
-            )
-
-    warn = []
-    metadata_factors = {}
-    for i, name in enumerate(names):
-        if name == CLASS_LABEL:
-            continue
-        if continuous_factor_bincounts and name in continuous_factor_bincounts:
-            metadata_factors[name] = digitize_factor_bins(data[:, i], continuous_factor_bincounts[name], name)
-        elif not is_categorical[i]:
-            warn.append(name)
-            metadata_factors[name] = data[:, i]
-        else:
-            metadata_factors[name] = data[:, i]
-
-    if warn:
-        warnings.warn(
-            f"The following factors appear to be continuous but did not have the desired number of bins specified: \n\
-            {warn}",
-            UserWarning,
-        )
-
-    return metadata_factors
-
-
 def normalize_expected_dist(expected_dist: NDArray[Any], observed_dist: NDArray[Any]) -> NDArray[Any]:
     """
     Normalize the expected label distribution to match the total number of labels in the observed distribution.
@@ -295,31 +205,19 @@ def label_parity(
 
 
 @set_metadata()
-def parity(
-    class_labels: ArrayLike,
-    metadata: Mapping[str, ArrayLike],
-    continuous_factor_bincounts: Mapping[str, int] | None = None,
-) -> ParityOutput[NDArray[np.float64]]:
+def parity(metadata: MetadataOutput) -> ParityOutput[NDArray[np.float64]]:
     """
-    Calculate chi-square statistics to assess the relationship between multiple factors
+    Calculate chi-square statistics to assess the linear relationship between multiple factors
     and class labels.
 
     This function computes the chi-square statistic for each metadata factor to determine if there is
-    a significant relationship between the factor values and class labels. The function handles both categorical
-    and discretized continuous factors.
+    a significant relationship between the factor values and class labels. The chi-square statistic is
+    only valid for linear relationships. If non-linear relationships exist, use `balance`.
 
     Parameters
     ----------
-    class_labels : ArrayLike
-        List of class labels for each image
-    metadata : Mapping[str, ArrayLike]
-        The dataset factors, which are per-image metadata attributes.
-        Each key of dataset_factors is a factor, whose value is the per-image factor values.
-    continuous_factor_bincounts : Mapping[str, int] or None, default None
-        A dictionary specifying the number of bins for discretizing the continuous factors.
-        The keys should correspond to the names of continuous factors in `metadata`,
-        and the values should be the number of bins to use for discretization.
-        If not provided, no discretization is applied.
+    metadata : MetadataOutput
+        Output after running `metadata_preprocessing`
 
     Returns
     -------
@@ -333,74 +231,66 @@ def parity(
     Warning
         If any cell in the contingency matrix has a value between 0 and 5, a warning is issued because this can
         lead to inaccurate chi-square calculations. It is recommended to ensure that each label co-occurs with
-        factor values either 0 times or at least 5 times. Alternatively, continuous-valued factors can be digitized
-        into fewer bins.
+        factor values either 0 times or at least 5 times.
 
     Note
     ----
-    - Each key of the ``continuous_factor_bincounts`` dictionary must occur as a key in data_factors.
     - A high score with a low p-value suggests that a metadata factor is strongly correlated with a class label.
     - The function creates a contingency matrix for each factor, where each entry represents the frequency of a
       specific factor value co-occurring with a particular class label.
     - Rows containing only zeros in the contingency matrix are removed before performing the chi-square test
       to prevent errors in the calculation.
 
+    See Also
+    --------
+    balance
+
     Examples
     --------
     Randomly creating some "continuous" and categorical variables using ``np.random.default_rng``
 
     >>> labels = np_random_gen.choice([0, 1, 2], (100))
-    >>> metadata = {
-    ...     "age": np_random_gen.choice([25, 30, 35, 45], (100)),
-    ...     "income": np_random_gen.choice([50000, 65000, 80000], (100)),
-    ...     "gender": np_random_gen.choice(["M", "F"], (100)),
-    ... }
+    >>> metadata_dict = [
+    ...     {
+    ...         "age": list(np_random_gen.choice([25, 30, 35, 45], (100))),
+    ...         "income": list(np_random_gen.choice([50000, 65000, 80000], (100))),
+    ...         "gender": list(np_random_gen.choice(["M", "F"], (100))),
+    ...     }
+    ... ]
     >>> continuous_factor_bincounts = {"age": 4, "income": 3}
-    >>> parity(labels, metadata, continuous_factor_bincounts)
+    >>> metadata = metadata_preprocessing(metadata_dict, labels, continuous_factor_bincounts)
+    >>> parity(metadata)
     ParityOutput(score=array([7.35731943, 5.46711299, 0.51506212]), p_value=array([0.28906231, 0.24263543, 0.77295762]), metadata_names=['age', 'income', 'gender'])
     """  # noqa: E501
-    if len(np.shape(class_labels)) > 1:
-        raise ValueError(
-            f"Got class labels with {len(np.shape(class_labels))}-dimensional",
-            f" shape {np.shape(class_labels)}, but expected a 1-dimensional array.",
-        )
-
-    data, names, is_categorical, _ = preprocess_metadata(class_labels, metadata)
-
-    factors = format_discretize_factors(data, names, is_categorical, continuous_factor_bincounts)
-
-    # unique class labels
-    class_idx = names.index(CLASS_LABEL)
-    u_cls = np.unique(data[:, class_idx])
-
-    chi_scores = np.zeros(len(factors))
-    p_values = np.zeros(len(factors))
+    chi_scores = np.zeros(metadata.discrete_data.shape[1])
+    p_values = np.zeros_like(chi_scores)
     not_enough_data = {}
-    for i, (current_factor_name, factor_values) in enumerate(factors.items()):
-        unique_factor_values = np.unique(factor_values)
-        contingency_matrix = np.zeros((len(unique_factor_values), u_cls.size))
+    for i, col_data in enumerate(metadata.discrete_data.T):
         # Builds a contingency matrix where entry at index (r,c) represents
         # the frequency of current_factor_name achieving value unique_factor_values[r]
         # at a data point with class c.
-
-        # TODO: Vectorize this nested for loop
-        for fi, factor_value in enumerate(unique_factor_values):
-            for label in u_cls:
-                with_both = np.bitwise_and((data[:, class_idx] == label), factor_values == factor_value)
-                contingency_matrix[fi, label] = np.sum(with_both)
-                if 0 < contingency_matrix[fi, label] < 5:
-                    if current_factor_name not in not_enough_data:
-                        not_enough_data[current_factor_name] = {}
-                    if factor_value not in not_enough_data[current_factor_name]:
-                        not_enough_data[current_factor_name][factor_value] = []
-                    not_enough_data[current_factor_name][factor_value].append(
-                        (label, int(contingency_matrix[fi, label]))
-                    )
+        results = crosstab(col_data, metadata.class_labels)
+        contingency_matrix = as_numpy(results.count)  # type: ignore
+
+        # Determines if any frequencies are too low
+        counts = np.nonzero(contingency_matrix < 5)
+        unique_factor_values = np.unique(col_data)
+        current_factor_name = metadata.discrete_factor_names[i]
+        for int_factor, int_class in zip(counts[0], counts[1]):
+            if contingency_matrix[int_factor, int_class] > 0:
+                factor_category = unique_factor_values[int_factor]
+                if current_factor_name not in not_enough_data:
+                    not_enough_data[current_factor_name] = {}
+                if factor_category not in not_enough_data[current_factor_name]:
+                    not_enough_data[current_factor_name][factor_category] = []
+                not_enough_data[current_factor_name][factor_category].append(
+                    (metadata.class_names[int_class], int(contingency_matrix[int_factor, int_class]))
+                )
 
         # This deletes rows containing only zeros,
         # because scipy.stats.chi2_contingency fails when there are rows containing only zeros.
         rowsums = np.sum(contingency_matrix, axis=1)
-        rowmask = np.where(rowsums)
+        rowmask = np.nonzero(rowsums)[0]
         contingency_matrix = contingency_matrix[rowmask]
 
         chi2, p, _, _ = chi2_contingency(contingency_matrix)
@@ -428,4 +318,4 @@ def parity(
             UserWarning,
         )
 
-    return ParityOutput(chi_scores, p_values, list(metadata.keys()))
+    return ParityOutput(chi_scores, p_values, metadata.discrete_factor_names)

dataeval/utils/gmm.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+TGMMData = TypeVar("TGMMData")
+
+
+@dataclass
+class GaussianMixtureModelParams(Generic[TGMMData]):
+    """
+    phi : TGMMData
+        Mixture component distribution weights.
+    mu : TGMMData
+        Mixture means.
+    cov : TGMMData
+        Mixture covariance.
+    L : TGMMData
+        Cholesky decomposition of `cov`.
+    log_det_cov : TGMMData
+        Log of the determinant of `cov`.
+    """
+
+    phi: TGMMData
+    mu: TGMMData
+    cov: TGMMData
+    L: TGMMData
+    log_det_cov: TGMMData

dataeval/utils/metadata.py

@@ -131,7 +131,9 @@ def _flatten_dict_inner(
     return items, size
 
 
-def _flatten_dict(d: Mapping[str, Any], sep: str, ignore_lists: bool, fully_qualified: bool) -> dict[str, Any]:
+def _flatten_dict(
+    d: Mapping[str, Any], sep: str, ignore_lists: bool, fully_qualified: bool
+) -> tuple[dict[str, Any], int]:
     """
     Flattens a dictionary and converts values to numeric values when possible.
 
@@ -165,7 +167,7 @@ def _flatten_dict(d: Mapping[str, Any], sep: str, ignore_lists: bool, fully_qual
             output[k] = cv
         elif not isinstance(cv, list):
             output[k] = cv if not size else [cv] * size
-    return output
+    return output, size if size is not None else 1
 
 
 def _is_metadata_dict_of_dicts(metadata: Mapping) -> bool:
@@ -188,7 +190,7 @@ def merge_metadata(
     ignore_lists: bool = False,
     fully_qualified: bool = False,
     as_numpy: bool = False,
-) -> dict[str, list[Any]] | dict[str, NDArray[Any]]:
+) -> tuple[dict[str, list[Any]] | dict[str, NDArray[Any]], NDArray[np.int_]]:
     """
     Merges a collection of metadata dictionaries into a single flattened dictionary of keys and values.
 
@@ -208,8 +210,10 @@ def merge_metadata(
 
     Returns
     -------
-    dict[str, list[Any]] | dict[str, NDArray[Any]]
+    dict[str, list[Any]] or dict[str, NDArray[Any]]
         A single dictionary containing the flattened data as lists or NumPy arrays
+    NDArray[np.int_]
+        Array defining where individual images start, helpful when working with object detection metadata
 
     Note
     ----
@@ -217,9 +221,12 @@ def merge_metadata(
 
     Example
     -------
-    >>> list_metadata = [{"common": 1, "target": [{"a": 1, "b": 3}, {"a": 2, "b": 4}], "source": "example"}]
-    >>> merge_metadata(list_metadata)
+    >>> list_metadata = [{"common": 1, "target": [{"a": 1, "b": 3, "c": 5}, {"a": 2, "b": 4}], "source": "example"}]
+    >>> reorganized_metadata, image_indicies = merge_metadata(list_metadata)
+    >>> reorganized_metadata
     {'common': [1, 1], 'a': [1, 2], 'b': [3, 4], 'source': ['example', 'example']}
+    >>> image_indicies
+    array([0])
     """
     merged: dict[str, list[Any]] = {}
     isect: set[str] = set()
@@ -236,8 +243,11 @@ def merge_metadata(
     else:
         dicts = list(metadata)
 
-    for d in dicts:
-        flattened = _flatten_dict(d, sep="_", ignore_lists=ignore_lists, fully_qualified=fully_qualified)
+    image_repeats = np.zeros(len(dicts))
+    for i, d in enumerate(dicts):
+        flattened, image_repeats[i] = _flatten_dict(
+            d, sep="_", ignore_lists=ignore_lists, fully_qualified=fully_qualified
+        )
         isect = isect.intersection(flattened.keys()) if isect else set(flattened.keys())
         union = union.union(flattened.keys())
         for k, v in flattened.items():
@@ -248,6 +258,16 @@ def merge_metadata(
 
     output: dict[str, Any] = {}
 
+    if image_repeats.sum() == image_repeats.size:
+        image_indicies = np.arange(image_repeats.size)
+    else:
+        image_ids = np.arange(image_repeats.size)
+        image_data = np.concatenate(
+            [np.repeat(image_ids[i], image_repeats[i]) for i in range(image_ids.size)], dtype=np.int_
+        )
+        _, image_unsorted = np.unique(image_data, return_index=True)
+        image_indicies = np.sort(image_unsorted)
+
     if keys:
         output["keys"] = np.array(keys) if as_numpy else keys
 
@@ -255,4 +275,4 @@ def merge_metadata(
         cv = _convert_type(merged[k])
         output[k] = np.array(cv) if as_numpy else cv
 
-    return output
+    return output, image_indicies

dataeval/utils/tensorflow/_internal/gmm.py

@@ -8,10 +8,11 @@ Licensed under Apache Software License (Apache 2.0)
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING
 
 import numpy as np
 
+from dataeval.utils.gmm import GaussianMixtureModelParams
 from dataeval.utils.lazy import lazyload
 
 if TYPE_CHECKING:
@@ -20,28 +21,7 @@ else:
     tf = lazyload("tensorflow")
 
 
-class GaussianMixtureModelParams(NamedTuple):
-    """
-    phi : tf.Tensor
-        Mixture component distribution weights.
-    mu : tf.Tensor
-        Mixture means.
-    cov : tf.Tensor
-        Mixture covariance.
-    L : tf.Tensor
-        Cholesky decomposition of `cov`.
-    log_det_cov : tf.Tensor
-        Log of the determinant of `cov`.
-    """
-
-    phi: tf.Tensor
-    mu: tf.Tensor
-    cov: tf.Tensor
-    L: tf.Tensor
-    log_det_cov: tf.Tensor
-
-
-def gmm_params(z: tf.Tensor, gamma: tf.Tensor) -> GaussianMixtureModelParams:
+def gmm_params(z: tf.Tensor, gamma: tf.Tensor) -> GaussianMixtureModelParams[tf.Tensor]:
     """
     Compute parameters of Gaussian Mixture Model.
 
@@ -81,7 +61,7 @@ def gmm_params(z: tf.Tensor, gamma: tf.Tensor) -> GaussianMixtureModelParams:
 
 def gmm_energy(
     z: tf.Tensor,
-    params: GaussianMixtureModelParams,
+    params: GaussianMixtureModelParams[tf.Tensor],
     return_mean: bool = True,
 ) -> tuple[tf.Tensor, tf.Tensor]:
     """