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

dataeval/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.72.2"
+__version__ = "0.73.1"
 
 from importlib.util import find_spec
 
@@ -12,12 +12,12 @@ from dataeval import detectors, metrics  # noqa: E402
 
 __all__ = ["detectors", "metrics"]
 
-if _IS_TORCH_AVAILABLE:  # pragma: no cover
+if _IS_TORCH_AVAILABLE:
     from dataeval import workflows
 
     __all__ += ["workflows"]
 
-if _IS_TENSORFLOW_AVAILABLE or _IS_TORCH_AVAILABLE:  # pragma: no cover
+if _IS_TENSORFLOW_AVAILABLE or _IS_TORCH_AVAILABLE:
     from dataeval import utils
 
     __all__ += ["utils"]

dataeval/detectors/__init__.py

@@ -7,7 +7,7 @@ from dataeval.detectors import drift, linters
 
 __all__ = ["drift", "linters"]
 
-if _IS_TENSORFLOW_AVAILABLE:  # pragma: no cover
+if _IS_TENSORFLOW_AVAILABLE:
     from dataeval.detectors import ood
 
     __all__ += ["ood"]

dataeval/detectors/drift/__init__.py

@@ -10,7 +10,7 @@ from dataeval.detectors.drift.ks import DriftKS
 
 __all__ = ["DriftCVM", "DriftKS", "DriftOutput", "updates"]
 
-if _IS_TORCH_AVAILABLE:  # pragma: no cover
+if _IS_TORCH_AVAILABLE:
     from dataeval.detectors.drift.mmd import DriftMMD, DriftMMDOutput
     from dataeval.detectors.drift.torch import preprocess_drift
     from dataeval.detectors.drift.uncertainty import DriftUncertainty

dataeval/detectors/drift/base.py

@@ -18,7 +18,7 @@ from typing import Any, Callable, Literal, TypeVar
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
-from dataeval.interop import as_numpy, to_numpy
+from dataeval.interop import as_numpy
 from dataeval.output import OutputMetadata, set_metadata
 
 R = TypeVar("R")
@@ -196,7 +196,7 @@ class BaseDrift:
         if correction not in ["bonferroni", "fdr"]:
             raise ValueError("`correction` must be `bonferroni` or `fdr`.")
 
-        self._x_ref = to_numpy(x_ref)
+        self._x_ref = as_numpy(x_ref)
         self.x_ref_preprocessed: bool = x_ref_preprocessed
 
         # Other attributes

dataeval/detectors/linters/clusterer.py

@@ -480,7 +480,7 @@ class Clusterer:
             samples = self.clusters[level][cluster_id].samples
             if len(samples) >= self._min_num_samples_per_cluster:
                 duplicates_std.append(self.clusters[level][cluster_id].dist_std)
-        diag_mask = np.ones_like(self._sqdmat, dtype=bool)
+        diag_mask = np.ones_like(self._sqdmat, dtype=np.bool_)
         np.fill_diagonal(diag_mask, 0)
         diag_mask = np.triu(diag_mask)
 

dataeval/detectors/ood/__init__.py

@@ -4,7 +4,7 @@ Out-of-distribution (OOD)` detectors identify data that is different from the da
 
 from dataeval import _IS_TENSORFLOW_AVAILABLE
 
-if _IS_TENSORFLOW_AVAILABLE:  # pragma: no cover
+if _IS_TENSORFLOW_AVAILABLE:
     from dataeval.detectors.ood.ae import OOD_AE
     from dataeval.detectors.ood.aegmm import OOD_AEGMM
     from dataeval.detectors.ood.base import OODOutput, OODScoreOutput

dataeval/detectors/ood/ae.py

@@ -10,18 +10,26 @@ from __future__ import annotations
 
 __all__ = ["OOD_AE"]
 
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
-import tensorflow as tf
-import tf_keras as keras
 from numpy.typing import ArrayLike
 
 from dataeval.detectors.ood.base import OODBase, OODScoreOutput
 from dataeval.interop import as_numpy
-from dataeval.utils.tensorflow._internal.autoencoder import AE
+from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.utils import predict_batch
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    import tf_keras as keras
+
+    import dataeval.utils.tensorflow._internal.models as tf_models
+else:
+    tf = lazyload("tensorflow")
+    keras = lazyload("tf_keras")
+    tf_models = lazyload("dataeval.utils.tensorflow._internal.models")
+
 
 class OOD_AE(OODBase):
     """
@@ -33,7 +41,7 @@ class OOD_AE(OODBase):
        An :term:`autoencoder<Autoencoder>` model.
     """
 
-    def __init__(self, model: AE) -> None:
+    def __init__(self, model: tf_models.AE) -> None:
         super().__init__(model)
 
     def fit(
@@ -41,7 +49,7 @@ class OOD_AE(OODBase):
         x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable[..., tf.Tensor] | None = None,
-        optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
+        optimizer: keras.optimizers.Optimizer | None = None,
         epochs: int = 20,
         batch_size: int = 64,
         verbose: bool = True,

dataeval/detectors/ood/aegmm.py

@@ -10,19 +10,27 @@ from __future__ import annotations
 
 __all__ = ["OOD_AEGMM"]
 
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
-import tensorflow as tf
-import tf_keras as keras
 from numpy.typing import ArrayLike
 
 from dataeval.detectors.ood.base import OODGMMBase, OODScoreOutput
 from dataeval.interop import to_numpy
-from dataeval.utils.tensorflow._internal.autoencoder import AEGMM
+from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.gmm import gmm_energy
 from dataeval.utils.tensorflow._internal.loss import LossGMM
 from dataeval.utils.tensorflow._internal.utils import predict_batch
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    import tf_keras as keras
+
+    import dataeval.utils.tensorflow._internal.models as tf_models
+else:
+    tf = lazyload("tensorflow")
+    keras = lazyload("tf_keras")
+    tf_models = lazyload("dataeval.utils.tensorflow._internal.models")
+
 
 class OOD_AEGMM(OODGMMBase):
     """
@@ -34,7 +42,7 @@ class OOD_AEGMM(OODGMMBase):
        An AEGMM model.
     """
 
-    def __init__(self, model: AEGMM) -> None:
+    def __init__(self, model: tf_models.AEGMM) -> None:
         super().__init__(model)
 
     def fit(
@@ -42,7 +50,7 @@ class OOD_AEGMM(OODGMMBase):
         x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable[..., tf.Tensor] | None = None,
-        optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
+        optimizer: keras.optimizers.Optimizer | None = None,
         epochs: int = 20,
         batch_size: int = 64,
         verbose: bool = True,

dataeval/detectors/ood/base.py

@@ -12,18 +12,24 @@ __all__ = ["OODOutput", "OODScoreOutput"]
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Callable, Literal, cast
+from typing import TYPE_CHECKING, Callable, Literal, cast
 
 import numpy as np
-import tensorflow as tf
-import tf_keras as keras
 from numpy.typing import ArrayLike, NDArray
 
 from dataeval.interop import to_numpy
 from dataeval.output import OutputMetadata, set_metadata
+from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.gmm import GaussianMixtureModelParams, gmm_params
 from dataeval.utils.tensorflow._internal.trainer import trainer
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    import tf_keras as keras
+else:
+    tf = lazyload("tensorflow")
+    keras = lazyload("tf_keras")
+
 
 @dataclass(frozen=True)
 class OODOutput(OutputMetadata):

dataeval/detectors/ood/llr.py

@@ -11,25 +11,31 @@ from __future__ import annotations
 __all__ = ["OOD_LLR"]
 
 from functools import partial
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
-import tensorflow as tf
-import tf_keras as keras
 from numpy.typing import ArrayLike, NDArray
-from tf_keras.layers import Input
-from tf_keras.models import Model
 
 from dataeval.detectors.ood.base import OODBase, OODScoreOutput
 from dataeval.interop import to_numpy
-from dataeval.utils.tensorflow._internal.pixelcnn import PixelCNN
+from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.trainer import trainer
 from dataeval.utils.tensorflow._internal.utils import predict_batch
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    import tf_keras as keras
+
+    import dataeval.utils.tensorflow._internal.models as tf_models
+else:
+    tf = lazyload("tensorflow")
+    keras = lazyload("tf_keras")
+    tf_models = lazyload("dataeval.utils.tensorflow._internal.models")
+
 
 def _build_model(
-    dist: PixelCNN, input_shape: tuple | None = None, filepath: str | None = None
-) -> tuple[keras.Model, PixelCNN]:
+    dist: tf_models.PixelCNN, input_shape: tuple | None = None, filepath: str | None = None
+) -> tuple[keras.Model, tf_models.PixelCNN]:
     """
     Create keras.Model from TF distribution.
 
@@ -46,9 +52,9 @@ def _build_model(
     -------
     TensorFlow model.
     """
-    x_in = Input(shape=input_shape)
+    x_in = keras.layers.Input(shape=input_shape)
     log_prob = dist.log_prob(x_in)
-    model = Model(inputs=x_in, outputs=log_prob)
+    model = keras.models.Model(inputs=x_in, outputs=log_prob)
     model.add_loss(-tf.reduce_mean(log_prob))
     if isinstance(filepath, str):
         model.load_weights(filepath)
@@ -109,13 +115,13 @@ class OOD_LLR(OODBase):
 
     def __init__(
         self,
-        model: PixelCNN,
-        model_background: PixelCNN | None = None,
+        model: tf_models.PixelCNN,
+        model_background: tf_models.PixelCNN | None = None,
         log_prob: Callable | None = None,
         sequential: bool = False,
     ) -> None:
-        self.dist_s: PixelCNN = model
-        self.dist_b: PixelCNN = (
+        self.dist_s: tf_models.PixelCNN = model
+        self.dist_b: tf_models.PixelCNN = (
             model.copy()
             if hasattr(model, "copy")
             else keras.models.clone_model(model)
@@ -135,7 +141,7 @@ class OOD_LLR(OODBase):
         x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable | None = None,
-        optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
+        optimizer: keras.optimizers.Optimizer | None = None,
         epochs: int = 20,
         batch_size: int = 64,
         verbose: bool = True,
@@ -176,7 +182,7 @@ class OOD_LLR(OODBase):
         """
         x_ref = to_numpy(x_ref)
         input_shape = x_ref.shape[1:]
-        optimizer = optimizer() if isinstance(optimizer, type) else optimizer
+        optimizer = keras.optimizers.Adam() if optimizer is None else optimizer
         # Separate into two separate optimizers, one for semantic model and one for background model
         optimizer_s = optimizer
         optimizer_b = optimizer.__class__.from_config(optimizer.get_config())

dataeval/detectors/ood/vae.py

@@ -10,19 +10,27 @@ from __future__ import annotations
 
 __all__ = ["OOD_VAE"]
 
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
-import tensorflow as tf
-import tf_keras as keras
 from numpy.typing import ArrayLike
 
 from dataeval.detectors.ood.base import OODBase, OODScoreOutput
 from dataeval.interop import to_numpy
-from dataeval.utils.tensorflow._internal.autoencoder import VAE
+from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.loss import Elbo
 from dataeval.utils.tensorflow._internal.utils import predict_batch
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    import tf_keras as keras
+
+    import dataeval.utils.tensorflow._internal.models as tf_models
+else:
+    tf = lazyload("tensorflow")
+    keras = lazyload("tf_keras")
+    tf_models = lazyload("dataeval.utils.tensorflow._internal.models")
+
 
 class OOD_VAE(OODBase):
     """
@@ -51,7 +59,7 @@ class OOD_VAE(OODBase):
     >>> result = metric.predict(dataset, ood_type="feature")
     """
 
-    def __init__(self, model: VAE, samples: int = 10) -> None:
+    def __init__(self, model: tf_models.VAE, samples: int = 10) -> None:
         super().__init__(model)
         self.samples = samples
 
@@ -60,7 +68,7 @@ class OOD_VAE(OODBase):
         x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable[..., tf.Tensor] = Elbo(0.05),
-        optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
+        optimizer: keras.optimizers.Optimizer | None = None,
         epochs: int = 20,
         batch_size: int = 64,
         verbose: bool = True,

dataeval/detectors/ood/vaegmm.py

@@ -10,20 +10,28 @@ from __future__ import annotations
 
 __all__ = ["OOD_VAEGMM"]
 
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
-import tensorflow as tf
-import tf_keras as keras
 from numpy.typing import ArrayLike
 
 from dataeval.detectors.ood.base import OODGMMBase, OODScoreOutput
 from dataeval.interop import to_numpy
-from dataeval.utils.tensorflow._internal.autoencoder import VAEGMM
+from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.gmm import gmm_energy
 from dataeval.utils.tensorflow._internal.loss import Elbo, LossGMM
 from dataeval.utils.tensorflow._internal.utils import predict_batch
 
+if TYPE_CHECKING:
+    import tensorflow as tf
+    import tf_keras as keras
+
+    import dataeval.utils.tensorflow._internal.models as tf_models
+else:
+    tf = lazyload("tensorflow")
+    keras = lazyload("tf_keras")
+    tf_models = lazyload("dataeval.utils.tensorflow._internal.models")
+
 
 class OOD_VAEGMM(OODGMMBase):
     """
@@ -37,7 +45,7 @@ class OOD_VAEGMM(OODGMMBase):
         Number of samples sampled to evaluate each instance.
     """
 
-    def __init__(self, model: VAEGMM, samples: int = 10) -> None:
+    def __init__(self, model: tf_models.VAEGMM, samples: int = 10) -> None:
         super().__init__(model)
         self.samples = samples
 
@@ -46,7 +54,7 @@ class OOD_VAEGMM(OODGMMBase):
         x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable[..., tf.Tensor] = LossGMM(elbo=Elbo(0.05)),
-        optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
+        optimizer: keras.optimizers.Optimizer | None = None,
         epochs: int = 20,
         batch_size: int = 64,
         verbose: bool = True,

dataeval/interop.py

@@ -37,13 +37,15 @@ def to_numpy(array: ArrayLike | None, copy: bool = True) -> NDArray[Any]:
     if isinstance(array, np.ndarray):
         return array.copy() if copy else array
 
-    tf = _try_import("tensorflow")
-    if tf and tf.is_tensor(array):
-        return array.numpy().copy() if copy else array.numpy()  # type: ignore
-
-    torch = _try_import("torch")
-    if torch and isinstance(array, torch.Tensor):
-        return array.detach().cpu().numpy().copy() if copy else array.detach().cpu().numpy()  # type: ignore
+    if array.__class__.__module__.startswith("tensorflow"):
+        tf = _try_import("tensorflow")
+        if tf and tf.is_tensor(array):
+            return array.numpy().copy() if copy else array.numpy()  # type: ignore
+
+    if array.__class__.__module__.startswith("torch"):
+        torch = _try_import("torch")
+        if torch and isinstance(array, torch.Tensor):
+            return array.detach().cpu().numpy().copy() if copy else array.detach().cpu().numpy()  # type: ignore
 
     return np.array(array, copy=copy)
 

dataeval/metrics/bias/balance.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 __all__ = ["BalanceOutput", "balance"]
 
+import contextlib
 import warnings
 from dataclasses import dataclass
 from typing import Any, Mapping
@@ -10,9 +11,12 @@ import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
 
-from dataeval.metrics.bias.metadata import entropy, heatmap, preprocess_metadata
+from dataeval.metrics.bias.metadata import CLASS_LABEL, entropy, heatmap, preprocess_metadata
 from dataeval.output import OutputMetadata, set_metadata
 
+with contextlib.suppress(ImportError):
+    from matplotlib.figure import Figure
+
 
 @dataclass(frozen=True)
 class BalanceOutput(OutputMetadata):
@@ -27,45 +31,43 @@ class BalanceOutput(OutputMetadata):
         Estimate of inter/intra-factor mutual information
     classwise : NDArray[np.float64]
         Estimate of mutual information between metadata factors and individual class labels
-    class_list: NDArray[np.int64]
-        Class labels for each value in the dataset
-    metadata_names: list[str]
+    class_list : NDArray
+        Array of the class labels present in the dataset
+    metadata_names : list[str]
         Names of each metadata factor
     """
 
     balance: NDArray[np.float64]
     factors: NDArray[np.float64]
     classwise: NDArray[np.float64]
-
-    class_list: NDArray[np.int64]
+    class_list: NDArray[Any]
     metadata_names: list[str]
 
     def plot(
         self,
-        row_labels: NDArray[Any] | None = None,
-        col_labels: NDArray[Any] | None = None,
+        row_labels: list[Any] | NDArray[Any] | None = None,
+        col_labels: list[Any] | NDArray[Any] | None = None,
         plot_classwise: bool = False,
-    ) -> None:
+    ) -> Figure:
         """
         Plot a heatmap of balance information
 
         Parameters
         ----------
-        row_labels: NDArray | None, default None
-            Array containing the labels for rows in the histogram
-        col_labels: NDArray | None, default None
-            Array containing the labels for columns in the histogram
-        plot_classwise: bool, default False
+        row_labels : ArrayLike or None, default None
+            List/Array containing the labels for rows in the histogram
+        col_labels : ArrayLike or None, default None
+            List/Array containing the labels for columns in the histogram
+        plot_classwise : bool, default False
             Whether to plot per-class balance instead of global balance
-
         """
         if plot_classwise:
             if row_labels is None:
-                row_labels = np.unique(self.class_list)
+                row_labels = self.class_list
             if col_labels is None:
                 col_labels = np.concatenate((["class"], self.metadata_names))
 
-            heatmap(
+            fig = heatmap(
                 self.classwise,
                 row_labels,
                 col_labels,
@@ -74,6 +76,7 @@ class BalanceOutput(OutputMetadata):
                 cbarlabel="Normalized Mutual Information",
             )
         else:
+            # Combine balance and factors results
             data = np.concatenate([self.balance[np.newaxis, 1:], self.factors], axis=0)
             # Create a mask for the upper triangle of the symmetrical array, ignoring the diagonal
             mask = np.triu(data + 1, k=0) < 1
@@ -87,12 +90,9 @@ class BalanceOutput(OutputMetadata):
             if col_labels is None:
                 col_labels = heat_labels[1:]
 
-            heatmap(
-                heat_data,
-                row_labels,
-                col_labels,
-                cbarlabel="Normalized Mutual Information",
-            )
+            fig = heatmap(heat_data, row_labels, col_labels, cbarlabel="Normalized Mutual Information")
+
+        return fig
 
 
 def validate_num_neighbors(num_neighbors: int) -> int:
@@ -116,19 +116,29 @@ def validate_num_neighbors(num_neighbors: int) -> int:
 
 
 @set_metadata("dataeval.metrics")
-def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neighbors: int = 5) -> BalanceOutput:
+def balance(
+    class_labels: ArrayLike,
+    metadata: Mapping[str, ArrayLike],
+    num_neighbors: int = 5,
+    continuous_factor_bincounts: Mapping[str, int] | None = None,
+) -> BalanceOutput:
     """
     Mutual information (MI) between factors (class label, metadata, label/image properties)
 
     Parameters
     ----------
-    class_labels: ArrayLike
+    class_labels : ArrayLike
         List of class labels for each image
-    metadata: Mapping[str, ArrayLike]
+    metadata : Mapping[str, ArrayLike]
         Dict of lists of metadata factors for each image
-    num_neighbors: int, default 5
+    num_neighbors : int, default 5
         Number of nearest neighbors to use for computing MI between discrete
         and continuous variables.
+    continuous_factor_bincounts : Mapping[str, int] or None, default None
+        The factors in metadata 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 metadata.
 
     Returns
     -------
@@ -148,7 +158,7 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     -------
     Return balance (mutual information) of factors with class_labels
 
-    >>> bal = balance(class_labels, metadata)
+    >>> bal = balance(class_labels, metadata, continuous_factor_bincounts=continuous_factor_bincounts)
     >>> bal.balance
     array([0.99999822, 0.13363788, 0.04505382, 0.02994455])
 
@@ -165,6 +175,7 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     array([[0.99999822, 0.13363788, 0.        , 0.        ],
            [0.99999822, 0.13363788, 0.        , 0.        ]])
 
+
     See Also
     --------
     sklearn.feature_selection.mutual_info_classif
@@ -172,18 +183,15 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     sklearn.metrics.mutual_info_score
     """
     num_neighbors = validate_num_neighbors(num_neighbors)
-    data, names, is_categorical = preprocess_metadata(class_labels, metadata)
+    data, names, is_categorical, unique_labels = preprocess_metadata(class_labels, metadata)
     num_factors = len(names)
     mi = np.empty((num_factors, num_factors))
     mi[:] = np.nan
 
-    class_idx = names.index("class_label")
-    class_lbl = np.array(data[:, class_idx], dtype=int)
-
     for idx in range(num_factors):
-        tgt = data[:, idx].astype(int)
+        tgt = data[:, idx].astype(np.intp)
 
-        if is_categorical[idx]:
+        if continuous_factor_bincounts and names[idx] not in continuous_factor_bincounts:
             mi[idx, :] = mutual_info_classif(
                 data,
                 tgt,
@@ -200,7 +208,7 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
                 random_state=0,
             )
 
-    ent_all = entropy(data, names, is_categorical, normalized=False)
+    ent_all = entropy(data, names, continuous_factor_bincounts, normalized=False)
     norm_factor = 0.5 * np.add.outer(ent_all, ent_all) + 1e-6
     # in principle MI should be symmetric, but it is not in practice.
     nmi = 0.5 * (mi + mi.T) / norm_factor
@@ -208,9 +216,8 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     factors = nmi[1:, 1:]
 
     # unique class labels
-    class_idx = names.index("class_label")
-    class_data = data[:, class_idx].astype(int)
-    u_cls = np.unique(class_data)
+    class_idx = names.index(CLASS_LABEL)
+    u_cls = np.unique(data[:, class_idx])
     num_classes = len(u_cls)
 
     # assume class is a factor
@@ -218,12 +225,11 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     classwise_mi[:] = np.nan
 
     # categorical variables, excluding class label
-    cat_mask = np.concatenate((is_categorical[:class_idx], is_categorical[(class_idx + 1) :]), axis=0).astype(int)
+    cat_mask = np.concatenate((is_categorical[:class_idx], is_categorical[(class_idx + 1) :]), axis=0).astype(np.intp)
 
-    tgt_bin = np.stack([class_data == cls for cls in u_cls]).T.astype(int)
-    ent_tgt_bin = entropy(
-        tgt_bin, names=[str(idx) for idx in range(num_classes)], is_categorical=[True for idx in range(num_classes)]
-    )
+    tgt_bin = np.stack([data[:, class_idx] == cls for cls in u_cls]).T.astype(np.intp)
+    names = [str(idx) for idx in range(num_classes)]
+    ent_tgt_bin = entropy(tgt_bin, names, continuous_factor_bincounts)
 
     # classification MI for discrete/categorical features
     for idx in range(num_classes):
@@ -240,4 +246,4 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     norm_factor = 0.5 * np.add.outer(ent_tgt_bin, ent_all) + 1e-6
     classwise = classwise_mi / norm_factor
 
-    return BalanceOutput(balance, factors, classwise, class_lbl, list(metadata.keys()))
+    return BalanceOutput(balance, factors, classwise, unique_labels, list(metadata.keys()))