dataeval 0.76.1__tar.gz → 0.82.0__tar.gz

{dataeval-0.76.1 → dataeval-0.82.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dataeval
-Version: 0.76.1
+Version: 0.82.0
 Summary: DataEval provides a simple interface to characterize image data and its impact on model performance across classification and object-detection tasks
 Home-page: https://dataeval.ai/
 License: MIT
@@ -21,7 +21,10 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Topic :: Scientific/Engineering
 Provides-Extra: all
+Requires-Dist: defusedxml (>=0.7.1)
+Requires-Dist: fast_hdbscan (==0.2.0)
 Requires-Dist: matplotlib (>=3.7.1) ; extra == "all"
+Requires-Dist: numba (>=0.59.1)
 Requires-Dist: numpy (>=1.24.2)
 Requires-Dist: pandas (>=2.0) ; extra == "all"
 Requires-Dist: pillow (>=10.3.0)
@@ -71,7 +74,7 @@ DataEval is easy to install, supports a wide range of Python versions, and is
 compatible with many of the most popular packages in the scientific and T&E
 communities.
 
-DataEval also has native interopability between JATIC's suite of tools when
+DataEval also has native interoperability between JATIC's suite of tools when
 using MAITE-compliant datasets and models.
 <!-- end JATIC interop -->
 

{dataeval-0.76.1 → dataeval-0.82.0}/README.md

@@ -32,7 +32,7 @@ DataEval is easy to install, supports a wide range of Python versions, and is
 compatible with many of the most popular packages in the scientific and T&E
 communities.
 
-DataEval also has native interopability between JATIC's suite of tools when
+DataEval also has native interoperability between JATIC's suite of tools when
 using MAITE-compliant datasets and models.
 <!-- end JATIC interop -->
 

{dataeval-0.76.1 → dataeval-0.82.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dataeval"
-version = "0.76.1" # dynamic
+version = "0.82.0" # dynamic
 description = "DataEval provides a simple interface to characterize image data and its impact on model performance across classification and object-detection tasks"
 license = "MIT"
 readme = "README.md"
@@ -42,6 +42,9 @@ packages = [
 [tool.poetry.dependencies]
 # required
 python = ">=3.9,<3.13"
+defusedxml = {version = ">=0.7.1"}
+fast_hdbscan = {version = "0.2.0"}  # 0.2.1 hits a bug in condense_tree comparing float to none
+numba = {version = ">=0.59.1"}
 numpy = {version = ">=1.24.2"}
 pillow = {version = ">=10.3.0"}
 requests = {version = "*"}
@@ -88,7 +91,7 @@ certifi = {version = ">=2024.07.04"}
 enum_tools = {version = ">=0.12.0", extras = ["sphinx"]}
 ipykernel = {version = ">=6.26.0"}
 ipywidgets = {version = ">=8.1.1"}
-jinja2 = {version = ">=3.1.5"}
+jinja2 = {version = ">=3.1.6"}
 jupyter-client = {version = ">=8.6.0"}
 jupyter-cache = {version = "*"}
 myst-nb = {version = ">=1.0.0"}
@@ -129,6 +132,11 @@ reportMissingImports = false
 norecursedirs = ["prototype"]
 testpaths = ["tests"]
 addopts = ["--pythonwarnings=ignore::DeprecationWarning", "--verbose", "--durations=20", "--durations-min=1.0"]
+markers = [
+  "required: marks tests for required features",
+  "optional: marks tests for optional features",
+  "requires_all: marks tests that require the all extras",
+]
 
 [tool.coverage.run]
 source = ["src/dataeval"]
@@ -143,8 +151,9 @@ exclude_also = [
 ]
 include = ["*/src/dataeval/*"]
 omit = [
-  "*/torch/blocks.py",
-  "*/torch/utils.py",
+  "*/torch/_blocks.py",
+  "*/_clusterer.py",
+  "*/_fast_mst.py",
 ]
 fail_under = 90
 
@@ -178,6 +187,9 @@ per-file-ignores = { "*.ipynb" = ["E402"] }
 [tool.ruff.lint.isort]
 known-first-party = ["dataeval"]
 
+[tool.ruff.lint.flake8-builtins]
+builtins-strict-checking = false
+
 [tool.ruff.format]
 quote-style = "double"
 indent-style = "space"
@@ -187,7 +199,7 @@ docstring-code-format = true
 docstring-code-line-length = "dynamic"
 
 [tool.codespell]
-skip = './*env*,./prototype,./output,./docs/build,./docs/source/.jupyter_cache,CHANGELOG.md,poetry.lock,*.html'
+skip = './*env*,./prototype,./output,./docs/build,./docs/source/.jupyter_cache,CHANGELOG.md,poetry.lock,*.html,./docs/source/*/data'
 ignore-words-list = ["Hart"]
 
 [build-system]

{dataeval-0.76.1 → dataeval-0.82.0}/src/dataeval/__init__.py

@@ -7,12 +7,12 @@ shifts that impact performance of deployed models.
 
 from __future__ import annotations
 
-__all__ = ["detectors", "log", "metrics", "utils", "workflows"]
-__version__ = "0.76.1"
+__all__ = ["config", "detectors", "log", "metrics", "typing", "utils", "workflows"]
+__version__ = "0.82.0"
 
 import logging
 
-from dataeval import detectors, metrics, utils, workflows
+from dataeval import config, detectors, metrics, typing, utils, workflows
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())
 

dataeval-0.82.0/src/dataeval/config.py

@@ -0,0 +1,77 @@
+"""
+Global configuration settings for DataEval.
+"""
+
+from __future__ import annotations
+
+__all__ = ["get_device", "set_device", "get_max_processes", "set_max_processes"]
+
+import torch
+from torch import device
+
+_device: device | None = None
+_processes: int | None = None
+
+
+def set_device(device: str | device | int) -> None:
+    """
+    Sets the default device to use when executing against a PyTorch backend.
+
+    Parameters
+    ----------
+    device : str or int or `torch.device`
+        The default device to use. See `torch.device <https://pytorch.org/docs/stable/tensor_attributes.html#torch.device>`_
+        documentation for more information.
+    """
+    global _device
+    _device = torch.device(device)
+
+
+def get_device(override: str | device | int | None = None) -> torch.device:
+    """
+    Returns the PyTorch device to use.
+
+    Parameters
+    ----------
+    override : str or int or `torch.device` or None, default None
+        The user specified override if provided, otherwise returns the default device.
+
+    Returns
+    -------
+    `torch.device`
+    """
+    if override is None:
+        global _device
+        return torch.get_default_device() if _device is None else _device
+    else:
+        return torch.device(override)
+
+
+def set_max_processes(processes: int | None) -> None:
+    """
+    Sets the maximum number of worker processes to use when running tasks that support parallel processing.
+
+    Parameters
+    ----------
+    processes : int or None
+        The maximum number of worker processes to use, or None to use
+        `os.process_cpu_count <https://docs.python.org/3/library/os.html#os.process_cpu_count>`_
+        to determine the number of worker processes.
+    """
+    global _processes
+    _processes = processes
+
+
+def get_max_processes() -> int | None:
+    """
+    Returns the maximum number of worker processes to use when running tasks that support parallel processing.
+
+    Returns
+    -------
+    int or None
+        The maximum number of worker processes to use, or None to use
+        `os.process_cpu_count <https://docs.python.org/3/library/os.html#os.process_cpu_count>`_
+        to determine the number of worker processes.
+    """
+    global _processes
+    return _processes

{dataeval-0.76.1 → dataeval-0.82.0}/src/dataeval/detectors/__init__.py

@@ -4,4 +4,4 @@ Detectors can determine if a dataset or individual images in a dataset are indic
 
 __all__ = ["drift", "linters", "ood"]
 
-from dataeval.detectors import drift, linters, ood
+from . import drift, linters, ood

dataeval-0.82.0/src/dataeval/detectors/drift/__init__.py

@@ -0,0 +1,22 @@
+"""
+:term:`Drift` detectors identify if the statistical properties of the data has changed.
+"""
+
+__all__ = [
+    "DriftCVM",
+    "DriftKS",
+    "DriftMMD",
+    "DriftMMDOutput",
+    "DriftOutput",
+    "DriftUncertainty",
+    "preprocess_drift",
+    "updates",
+]
+
+from dataeval.detectors.drift import updates
+from dataeval.detectors.drift._cvm import DriftCVM
+from dataeval.detectors.drift._ks import DriftKS
+from dataeval.detectors.drift._mmd import DriftMMD
+from dataeval.detectors.drift._torch import preprocess_drift
+from dataeval.detectors.drift._uncertainty import DriftUncertainty
+from dataeval.outputs._drift import DriftMMDOutput, DriftOutput

dataeval-0.76.1/src/dataeval/detectors/drift/base.py → dataeval-0.82.0/src/dataeval/detectors/drift/_base.py

@@ -10,86 +10,29 @@ from __future__ import annotations
 
 __all__ = []
 
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
+import math
+from abc import abstractmethod
 from functools import wraps
-from typing import Any, Callable, Literal, TypeVar
+from typing import Any, Callable, Literal, Protocol, TypeVar, runtime_checkable
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import NDArray
 
-from dataeval.interop import as_numpy
-from dataeval.output import Output, set_metadata
+from dataeval.outputs import DriftOutput
+from dataeval.outputs._base import set_metadata
+from dataeval.typing import Array, ArrayLike
+from dataeval.utils._array import as_numpy, to_numpy
 
 R = TypeVar("R")
 
 
-class UpdateStrategy(ABC):
+@runtime_checkable
+class UpdateStrategy(Protocol):
     """
-    Updates reference dataset for drift detector
-
-    Parameters
-    ----------
-    n : int
-        Update with last n instances seen by the detector.
-    """
-
-    def __init__(self, n: int) -> None:
-        self.n = n
-
-    @abstractmethod
-    def __call__(self, x_ref: NDArray[Any], x: NDArray[Any], count: int) -> NDArray[Any]:
-        """Abstract implementation of update strategy"""
-
-
-@dataclass(frozen=True)
-class DriftBaseOutput(Output):
-    """
-    Base output class for Drift Detector classes
-
-    Attributes
-    ----------
-    is_drift : bool
-        Drift prediction for the images
-    threshold : float
-        Threshold after multivariate correction if needed
+    Protocol for reference dataset update strategy for drift detectors
     """
 
-    is_drift: bool
-    threshold: float
-    p_val: float
-    distance: float
-
-
-@dataclass(frozen=True)
-class DriftOutput(DriftBaseOutput):
-    """
-    Output class for :class:`DriftCVM`, :class:`DriftKS`, and :class:`DriftUncertainty` drift detectors.
-
-    Attributes
-    ----------
-    is_drift : bool
-        :term:`Drift` prediction for the images
-    threshold : float
-        Threshold after multivariate correction if needed
-    feature_drift : NDArray
-        Feature-level array of images detected to have drifted
-    feature_threshold : float
-        Feature-level threshold to determine drift
-    p_vals : NDArray
-        Feature-level p-values
-    distances : NDArray
-        Feature-level distances
-    """
-
-    # is_drift: bool
-    # threshold: float
-    # p_val: float
-    # distance: float
-    feature_drift: NDArray[np.bool_]
-    feature_threshold: float
-    p_vals: NDArray[np.float32]
-    distances: NDArray[np.float32]
+    def __call__(self, x_ref: NDArray[Any], x: NDArray[Any], count: int) -> NDArray[Any]: ...
 
 
 def update_x_ref(fn: Callable[..., R]) -> Callable[..., R]:
@@ -196,7 +139,7 @@ class BaseDrift:
         if correction not in ["bonferroni", "fdr"]:
             raise ValueError("`correction` must be `bonferroni` or `fdr`.")
 
-        self._x_ref = as_numpy(x_ref)
+        self._x_ref = x_ref
         self.x_ref_preprocessed: bool = x_ref_preprocessed
 
         # Other attributes
@@ -204,25 +147,25 @@ class BaseDrift:
         self.update_x_ref = update_x_ref
         self.preprocess_fn = preprocess_fn
         self.correction = correction
-        self.n: int = len(self._x_ref)
+        self.n: int = len(x_ref)
 
         # Ref counter for preprocessed x
         self._x_refcount = 0
 
     @property
-    def x_ref(self) -> NDArray[Any]:
+    def x_ref(self) -> ArrayLike:
         """
         Retrieve the reference data, applying preprocessing if not already done.
 
         Returns
         -------
-        NDArray
+        ArrayLike
             The reference dataset (`x_ref`), preprocessed if needed.
         """
         if not self.x_ref_preprocessed:
             self.x_ref_preprocessed = True
             if self.preprocess_fn is not None:
-                self._x_ref = as_numpy(self.preprocess_fn(self._x_ref))
+                self._x_ref = self.preprocess_fn(self._x_ref)
 
         return self._x_ref
 
@@ -323,32 +266,44 @@ class BaseDriftUnivariate(BaseDrift):
         # lazy process n_features as needed
         if not isinstance(self._n_features, int):
             # compute number of features for the univariate tests
-            if not isinstance(self.preprocess_fn, Callable) or self.x_ref_preprocessed:
-                # infer features from preprocessed reference data
-                self._n_features = self.x_ref.reshape(self.x_ref.shape[0], -1).shape[-1]
-            else:
-                # infer number of features after applying preprocessing step
-                x = as_numpy(self.preprocess_fn(self._x_ref[0:1]))  # type: ignore
-                self._n_features = x.reshape(x.shape[0], -1).shape[-1]
+            x_ref = (
+                self.x_ref
+                if self.preprocess_fn is None or self.x_ref_preprocessed
+                else self.preprocess_fn(self._x_ref[0:1])
+            )
+            # infer features from preprocessed reference data
+            shape = x_ref.shape if isinstance(x_ref, Array) else as_numpy(x_ref).shape
+            self._n_features = int(math.prod(shape[1:]))  # Multiplies all channel sizes after first
 
         return self._n_features
 
     @preprocess_x
-    @abstractmethod
     def score(self, x: ArrayLike) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
         """
-        Abstract method to calculate feature scores after preprocessing.
+        Calculates p-values and test statistics per feature.
 
         Parameters
         ----------
         x : ArrayLike
-            The batch of data to calculate univariate :term:`drift<Drift>` scores for each feature.
+            Batch of instances
 
         Returns
         -------
         tuple[NDArray, NDArray]
-            A tuple containing p-values and distance :term:`statistics<Statistics>` for each feature.
+            Feature level p-values and test statistics
         """
+        x_np = to_numpy(x)
+        x_np = x_np.reshape(x_np.shape[0], -1)
+        x_ref_np = as_numpy(self.x_ref)
+        x_ref_np = x_ref_np.reshape(x_ref_np.shape[0], -1)
+        p_val = np.zeros(self.n_features, dtype=np.float32)
+        dist = np.zeros_like(p_val)
+        for f in range(self.n_features):
+            dist[f], p_val[f] = self._score_fn(x_ref_np[:, f], x_np[:, f])
+        return p_val, dist
+
+    @abstractmethod
+    def _score_fn(self, x: NDArray[np.float32], y: NDArray[np.float32]) -> tuple[np.float32, np.float32]: ...
 
     def _apply_correction(self, p_vals: NDArray) -> tuple[bool, float]:
         """

dataeval-0.76.1/src/dataeval/detectors/drift/cvm.py → dataeval-0.82.0/src/dataeval/detectors/drift/_cvm.py

@@ -13,11 +13,11 @@ __all__ = []
 from typing import Callable, Literal
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import NDArray
 from scipy.stats import cramervonmises_2samp
 
-from dataeval.detectors.drift.base import BaseDriftUnivariate, UpdateStrategy, preprocess_x
-from dataeval.interop import to_numpy
+from dataeval.detectors.drift._base import BaseDriftUnivariate, UpdateStrategy
+from dataeval.typing import ArrayLike
 
 
 class DriftCVM(BaseDriftUnivariate):
@@ -55,6 +55,21 @@ class DriftCVM(BaseDriftUnivariate):
         Number of features used in the statistical test. No need to pass it if no
         preprocessing takes place. In case of a preprocessing step, this can also
         be inferred automatically but could be more expensive to compute.
+
+    Example
+    -------
+    >>> from functools import partial
+    >>> from dataeval.detectors.drift import preprocess_drift
+
+    Use a preprocess function to encode images before testing for drift
+
+    >>> preprocess_fn = partial(preprocess_drift, model=encoder, batch_size=64)
+    >>> drift = DriftCVM(train_images, preprocess_fn=preprocess_fn)
+
+    Test incoming images for drift
+
+    >>> drift.predict(test_images).drifted
+    True
     """
 
     def __init__(
@@ -77,28 +92,6 @@ class DriftCVM(BaseDriftUnivariate):
             n_features=n_features,
         )
 
-    @preprocess_x
-    def score(self, x: ArrayLike) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
-        """
-        Performs the two-sample Cramér-von Mises test(s), computing the :term:`p-value<P-value>` and
-        test statistic per feature.
-
-        Parameters
-        ----------
-        x : ArrayLike
-            Batch of instances.
-
-        Returns
-        -------
-        tuple[NDArray, NDArray]
-            Feature level p-values and CVM statistic
-        """
-        x_np = to_numpy(x)
-        x_np = x_np.reshape(x_np.shape[0], -1)
-        x_ref = self.x_ref.reshape(self.x_ref.shape[0], -1)
-        p_val = np.zeros(self.n_features, dtype=np.float32)
-        dist = np.zeros_like(p_val)
-        for f in range(self.n_features):
-            result = cramervonmises_2samp(x_ref[:, f], x_np[:, f], method="auto")
-            p_val[f], dist[f] = result.pvalue, result.statistic
-        return p_val, dist
+    def _score_fn(self, x: NDArray[np.float32], y: NDArray[np.float32]) -> tuple[np.float32, np.float32]:
+        result = cramervonmises_2samp(x, y, method="auto")
+        return np.float32(result.statistic), np.float32(result.pvalue)

dataeval-0.76.1/src/dataeval/detectors/drift/ks.py → dataeval-0.82.0/src/dataeval/detectors/drift/_ks.py

@@ -13,11 +13,11 @@ __all__ = []
 from typing import Callable, Literal
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import NDArray
 from scipy.stats import ks_2samp
 
-from dataeval.detectors.drift.base import BaseDriftUnivariate, UpdateStrategy, preprocess_x
-from dataeval.interop import to_numpy
+from dataeval.detectors.drift._base import BaseDriftUnivariate, UpdateStrategy
+from dataeval.typing import ArrayLike
 
 
 class DriftKS(BaseDriftUnivariate):
@@ -58,6 +58,21 @@ class DriftKS(BaseDriftUnivariate):
         Number of features used in the statistical test. No need to pass it if no
         preprocessing takes place. In case of a preprocessing step, this can also
         be inferred automatically but could be more expensive to compute.
+
+    Example
+    -------
+    >>> from functools import partial
+    >>> from dataeval.detectors.drift import preprocess_drift
+
+    Use a preprocess function to encode images before testing for drift
+
+    >>> preprocess_fn = partial(preprocess_drift, model=encoder, batch_size=64)
+    >>> drift = DriftKS(train_images, preprocess_fn=preprocess_fn)
+
+    Test incoming images for drift
+
+    >>> drift.predict(test_images).drifted
+    True
     """
 
     def __init__(
@@ -84,26 +99,5 @@ class DriftKS(BaseDriftUnivariate):
         # Other attributes
         self.alternative = alternative
 
-    @preprocess_x
-    def score(self, x: ArrayLike) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
-        """
-        Compute KS scores and :term:Statistics` per feature.
-
-        Parameters
-        ----------
-        x : ArrayLike
-            Batch of instances.
-
-        Returns
-        -------
-        tuple[NDArray, NDArray]
-            Feature level :term:p-values and KS statistic
-        """
-        x = to_numpy(x)
-        x = x.reshape(x.shape[0], -1)
-        x_ref = self.x_ref.reshape(self.x_ref.shape[0], -1)
-        p_val = np.zeros(self.n_features, dtype=np.float32)
-        dist = np.zeros_like(p_val)
-        for f in range(self.n_features):
-            dist[f], p_val[f] = ks_2samp(x_ref[:, f], x[:, f], alternative=self.alternative, method="exact")
-        return p_val, dist
+    def _score_fn(self, x: NDArray[np.float32], y: NDArray[np.float32]) -> tuple[np.float32, np.float32]:
+        return ks_2samp(x, y, alternative=self.alternative, method="exact")

dataeval-0.76.1/src/dataeval/detectors/drift/mmd.py → dataeval-0.82.0/src/dataeval/detectors/drift/_mmd.py

@@ -10,43 +10,16 @@ from __future__ import annotations
 
 __all__ = []
 
-from dataclasses import dataclass
 from typing import Callable
 
 import torch
-from numpy.typing import ArrayLike
 
-from dataeval.detectors.drift.base import BaseDrift, DriftBaseOutput, UpdateStrategy, preprocess_x, update_x_ref
-from dataeval.detectors.drift.torch import GaussianRBF, mmd2_from_kernel_matrix
-from dataeval.interop import as_numpy
-from dataeval.output import set_metadata
-from dataeval.utils.torch.internal import get_device
-
-
-@dataclass(frozen=True)
-class DriftMMDOutput(DriftBaseOutput):
-    """
-    Output class for :class:`DriftMMD` :term:`drift<Drift>` detector.
-
-    Attributes
-    ----------
-    is_drift : bool
-        Drift prediction for the images
-    threshold : float
-        :term:`P-Value` used for significance of the permutation test
-    p_val : float
-        P-value obtained from the permutation test
-    distance : float
-        MMD^2 between the reference and test set
-    distance_threshold : float
-        MMD^2 threshold above which drift is flagged
-    """
-
-    # is_drift: bool
-    # threshold: float
-    # p_val: float
-    # distance: float
-    distance_threshold: float
+from dataeval.config import get_device
+from dataeval.detectors.drift._base import BaseDrift, UpdateStrategy, preprocess_x, update_x_ref
+from dataeval.detectors.drift._torch import GaussianRBF, mmd2_from_kernel_matrix
+from dataeval.outputs import DriftMMDOutput
+from dataeval.outputs._base import set_metadata
+from dataeval.typing import ArrayLike
 
 
 class DriftMMD(BaseDrift):
@@ -84,6 +57,21 @@ class DriftMMD(BaseDrift):
     device : str | None, default None
         Device type used. The default None uses the GPU and falls back on CPU.
         Can be specified by passing either 'cuda', 'gpu' or 'cpu'.
+
+    Example
+    -------
+    >>> from functools import partial
+    >>> from dataeval.detectors.drift import preprocess_drift
+
+    Use a preprocess function to encode images before testing for drift
+
+    >>> preprocess_fn = partial(preprocess_drift, model=encoder, batch_size=64)
+    >>> drift = DriftMMD(train_images, preprocess_fn=preprocess_fn)
+
+    Test incoming images for drift
+
+    >>> drift.predict(test_images).drifted
+    True
     """
 
     def __init__(
@@ -110,12 +98,12 @@ class DriftMMD(BaseDrift):
         self.device: torch.device = get_device(device)
 
         # initialize kernel
-        sigma_tensor = torch.from_numpy(as_numpy(sigma)).to(self.device) if sigma is not None else None
+        sigma_tensor = torch.as_tensor(sigma, device=self.device) if sigma is not None else None
         self._kernel = GaussianRBF(sigma_tensor).to(self.device)
 
         # compute kernel matrix for the reference data
         if self._infer_sigma or isinstance(sigma_tensor, torch.Tensor):
-            x = torch.from_numpy(self.x_ref).to(self.device)
+            x = torch.as_tensor(self.x_ref, device=self.device)
             self._k_xx = self._kernel(x, x, infer_sigma=self._infer_sigma)
             self._infer_sigma = False
         else:
@@ -147,21 +135,21 @@ class DriftMMD(BaseDrift):
             p-value obtained from the permutation test, MMD^2 between the reference and test set,
             and MMD^2 threshold above which :term:`drift<Drift>` is flagged
         """
-        x = as_numpy(x)
-        x_ref = torch.from_numpy(self.x_ref).to(self.device)
-        n = x.shape[0]
-        kernel_mat = self._kernel_matrix(x_ref, torch.from_numpy(x).to(self.device))
+        x_ref = torch.as_tensor(self.x_ref, device=self.device)
+        x_test = torch.as_tensor(x, device=self.device)
+        n = x_test.shape[0]
+        kernel_mat = self._kernel_matrix(x_ref, x_test)
         kernel_mat = kernel_mat - torch.diag(kernel_mat.diag())  # zero diagonal
         mmd2 = mmd2_from_kernel_matrix(kernel_mat, n, permute=False, zero_diag=False)
-        mmd2_permuted = torch.Tensor(
-            [mmd2_from_kernel_matrix(kernel_mat, n, permute=True, zero_diag=False) for _ in range(self.n_permutations)]
+        mmd2_permuted = torch.tensor(
+            [mmd2_from_kernel_matrix(kernel_mat, n, permute=True, zero_diag=False)] * self.n_permutations,
+            device=self.device,
         )
-        mmd2, mmd2_permuted = mmd2.detach().cpu(), mmd2_permuted.detach().cpu()
         p_val = (mmd2 <= mmd2_permuted).float().mean()
         # compute distance threshold
         idx_threshold = int(self.p_val * len(mmd2_permuted))
         distance_threshold = torch.sort(mmd2_permuted, descending=True).values[idx_threshold]
-        return p_val.numpy().item(), mmd2.numpy().item(), distance_threshold.numpy().item()
+        return float(p_val.item()), float(mmd2.item()), float(distance_threshold.item())
 
     @set_metadata
     @preprocess_x