Repository: ViCCo-Group/thingsvision (language: python)

Issue:
Moving to torch 2.0 + 
Hi everyone
                   The library is amazing. Do you plan to migrate to PyTorch 2.0+ ?


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 3aecc90d..f1c1cb68 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@
         <img src="https://img.shields.io/pypi/dm/thingsvision" alt="downloads">
     </a>
     <a href="https://www.python.org/" rel="nofollow">
-        <img src="https://img.shields.io/badge/python-3.8%20%7C%203.9%20%7C%203.10-blue.svg" alt="Python version" />
+        <img src="https://img.shields.io/badge/python-3.9%20%7C%203.10-blue.svg" alt="Python version" />
     </a>
     <a href="https://github.com/ViCCo-Group/thingsvision/blob/master/LICENSE" rel="nofollow">
         <img src="https://img.shields.io/pypi/l/thingsvision" alt="License" />
diff --git a/docs/RSA.md b/docs/RSA.md
index 91f4791f..769be001 100644
--- a/docs/RSA.md
+++ b/docs/RSA.md
@@ -1,5 +1,5 @@
 ---
-title: RSA tools
+title: Comparing representations
 nav_order: 8
 ---
 
@@ -23,10 +23,15 @@ corr_coeff = correlate_rdms(rdm_dnn, rdm_human, correlation='pearson')
 Perform CKA to compare image features of two different model architectures for the same layer, or two different layers of the same architecture.
 
 ```python
-from thingsvision.core.cka import CKA
+from thingsvision.core.cka import get_cka
 
+backend = "torch" # can be set to either 'torch' or 'numpy'
 m = # number of images (e.g., features_i.shape[0])
-kernel = 'linear'
-cka = CKA(m=m, kernel=kernel)
+kernel = 'linear' # linear or rbf kernel (for rbf kernel define sigma, i.e., the width of the Gaussian)
+unbiased = True # whether to compute an unbiased version of CKA
+device = "cuda" # only necessary to be defined for the 'torch' backend (NumPy runs on CPU only)
+sigma = None # needs to be defined for 'rbf' kernel
+
+cka = get_cka(backend=backend, m=m, kernel=kernel, unbiased=unbiased, device=device, sigma=sigma)
 rho = cka.compare(X=features_i, Y=features_j)
 ```
diff --git a/docs/index.md b/docs/index.md
index 473b276f..1a6d64df 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -22,7 +22,8 @@ Neural networks come from different sources. With `thingsvision`, you can extrac
 - `ssl` (self-supervised learning models)
   - `simclr-rn50`, `mocov2-rn50`, `barlowtwins-rn50`, `pirl-rn50`
   - `jigsaw-rn50`, `rotnet-rn50`, `swav-rn50`, `vicreg-rn50`
-  - `dino-rn50`, `dino-xcit-{small/medium}-{12/24}-p{8/16}`, `dino-vit-{tiny/small/base}-p{8/16}`, `dinov2-vit-{small/base/large/giant}-p14`
+  - `dino-rn50`, `dino-xcit-{small/medium}-{12/24}-p{8/16}`
+  - `dino-vit-{tiny/small/base}-p{8/16}`, `dinov2-vit-{small/base/large/giant}-p14`
   - `mae-vit-{base/large}-p16`, `mae-vit-huge-p14`<br>
 - [OpenCLIP](https://github.com/mlfoundations/open_clip) models (CLIP trained on LAION-{400M/2B/5B})
 - [CLIP](https://github.com/openai/CLIP) models (CLIP trained on WiT)
diff --git a/requirements.txt b/requirements.txt
index 17d0ac71..8e22d1d3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,10 +13,10 @@ scipy
 tensorflow==2.9.* ; sys_platform != 'darwin' or platform_machine != 'arm64'
 tensorflow-macos==2.9.* ; sys_platform == 'darwin' and platform_machine == 'arm64'
 timm
-torch>=1.9.0
+torch>=2.0.0
 torchvision==0.15.2
 torchtyping
 tqdm
 dreamsim==0.1.3
 git+https://github.com/openai/CLIP.git
-git+https://github.com/serre-lab/Harmonization.git
+git+https://github.com/serre-lab/Harmonization.git
\ No newline at end of file
diff --git a/setup.py b/setup.py
index c767c8d0..e5ee44aa 100644
--- a/setup.py
+++ b/setup.py
@@ -21,7 +21,7 @@
     "tensorflow==2.9.* ; sys_platform != 'darwin' or platform_machine != 'arm64'",
     "tensorflow-macos==2.9.* ; sys_platform == 'darwin' and platform_machine == 'arm64'",
     "timm",
-    "torch>=1.9.0",
+    "torch>=2.0.0",
     "torchvision==0.15.2",
     "torchtyping",
     "tqdm",
diff --git a/thingsvision/_version.py b/thingsvision/_version.py
index f4a337e1..e5e59e38 100644
--- a/thingsvision/_version.py
+++ b/thingsvision/_version.py
@@ -1,1 +1,1 @@
-__version__ = "2.5.4"
+__version__ = "2.6.0"
diff --git a/thingsvision/core/cka/__init__.py b/thingsvision/core/cka/__init__.py
index f3804f24..b06040d0 100644
--- a/thingsvision/core/cka/__init__.py
+++ b/thingsvision/core/cka/__init__.py
@@ -1,1 +1,1 @@
-from .base import CKA
\ No newline at end of file
+from .helpers import get_cka
diff --git a/thingsvision/core/cka/base.py b/thingsvision/core/cka/base.py
deleted file mode 100644
index 5b6461da..00000000
--- a/thingsvision/core/cka/base.py
+++ /dev/null
@@ -1,65 +0,0 @@
-
-import math
-import numpy as np
-from dataclasses import dataclass
-
-Array = np.ndarray
-
-@dataclass
-class CKA:
-    m: int # number of examples
-    kernel: str
-    sigma: float=None
-        
-    def __post_init__(self) -> None:
-        self.H = self.centering_matrix(self.m)
-    
-    @staticmethod
-    def centering_matrix(m: int) -> Array:
-        """Compute the centering matrix H."""
-        H = np.eye(m) - np.ones((m, m)) / m
-        return H
-    
-    def centering(self, K: Array) -> Array:
-        """Centering of the gram matrix K."""
-        if not np.allclose(K, K.T):
-            raise ValueError('\nInput array must be a symmetric matrix.\n')
-        K_c = self.H @ K @ self.H
-        return K_c
-    
-    def apply_kernel(self, X: Array) -> Array:
-        """Compute the gram matrix K."""
-        try:
-            K = getattr(self, f'{self.kernel}_kernel')(X)
-        except AttributeError:
-            raise NotImplementedError
-        return K
-    
-    def linear_kernel(self, X: Array) -> Array:
-        return X @ X.T
-    
-    def rbf_kernel(self, X: Array, sigma=None) -> Array:
-        GX = X @ X.T
-        KX = np.diag(GX) - GX + (np.diag(GX) - GX).T
-        if sigma is None:
-            mdist = np.median(KX[KX != 0])
-            sigma = math.sqrt(mdist)
-        KX *= - 0.5 / sigma ** 2
-        KX = np.exp(KX)
-        return KX
-    
-    def hsic(self, X: Array, Y: Array) -> float:
-        K = self.apply_kernel(X)
-        L = self.apply_kernel(Y)
-        K_c = self.centering(K)
-        L_c = self.centering(L)
-        # np.sum(K_c * L_c) is equivalent to K_c.flatten() @ L_c.flatten() or in math
-        # sum_{i=0}^{m} sum_{j=0}^{m} K^{\prime}_{ij} * L^{\prime}_{ij} = vec(K_c)^{T}vec(L_c)
-        return np.sum(K_c * L_c)
-    
-    def compare(self, X: Array, Y: Array) -> float:
-        hsic_xy = self.hsic(X, Y)
-        hsic_xx = self.hsic(X, X)
-        hsic_yy = self.hsic(Y, Y)
-        rho = hsic_xy / np.sqrt(hsic_xx * hsic_yy)
-        return rho
\ No newline at end of file
diff --git a/thingsvision/core/cka/cka_base.py b/thingsvision/core/cka/cka_base.py
new file mode 100644
index 00000000..3798e1a6
--- /dev/null
+++ b/thingsvision/core/cka/cka_base.py
@@ -0,0 +1,62 @@
+import abc
+from typing import Optional, Union
+
+import numpy as np
+import torch
+
+Array = np.ndarray
+Tensor = torch.Tensor
+
+
+class CKABase(metaclass=abc.ABCMeta):
+    def __init__(
+        self, m: int, kernel: str, unbiased: bool = False, sigma: Optional[float] = 1.0
+    ) -> None:
+        """
+        CKA abstract base class from which other CKA classes inherit.
+        Args:
+            m (int) - number of images / examples in a mini-batch or the full dataset;
+            kernel (str) - 'linear' or 'rbf' kernel for computing the gram matrix;
+            unbiased (bool) - whether to compute an unbiased version of CKA;
+            sigma (float) - for 'rbf' kernel sigma defines the width of the Gaussian;
+        """
+        self.m = m  # number of examples
+        self.kernel = kernel  # linear or rbf kernel
+        self.unbiased = unbiased  # whether to use the unbiased version of CKA
+        self.sigma = sigma  # width of the Gaussian in an rbf kernel
+
+    @abc.abstractmethod
+    def centering(self, K: Union[Array, Tensor]) -> Union[Array, Tensor]:
+        """Centering of the (square) gram matrix K."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def apply_kernel(self, X: Union[Array, Tensor]) -> Union[Array, Tensor]:
+        """Compute the (square) gram matrix K."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def linear_kernel(self, X: Union[Array, Tensor]) -> Union[Array, Tensor]:
+        """Use a linear kernel for computing the gram matrix."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def rbf_kernel(
+        self, X: Union[Array, Tensor], sigma: float = None
+    ) -> Union[Array, Tensor]:
+        """Use an rbf kernel for computing the gram matrix. Sigma defines the width."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _hsic(
+        self, X: Union[Array, Tensor], Y: Union[Array, Tensor]
+    ) -> Union[Array, Tensor]:
+        """Compute the Hilbert-Schmidt independence criterion."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def compare(
+        self, X: Union[Array, Tensor], Y: Union[Array, Tensor]
+    ) -> Union[Array, Tensor]:
+        """Compare two representation spaces X and Y using CKA."""
+        raise NotImplementedError
diff --git a/thingsvision/core/cka/cka_numpy.py b/thingsvision/core/cka/cka_numpy.py
new file mode 100644
index 00000000..d276c802
--- /dev/null
+++ b/thingsvision/core/cka/cka_numpy.py
@@ -0,0 +1,89 @@
+import math
+from typing import Optional
+
+import numpy as np
+
+from .cka_base import CKABase
+
+Array = np.ndarray
+
+
+class CKANumPy(CKABase):
+    def __init__(
+        self, m: int, kernel: str, unbiased: bool = False, sigma: Optional[float] = 1.0
+    ) -> None:
+        """
+        NumPy implementation of CKA.
+        Args:
+            m (int) - number of images / examples in a mini-batch or the full dataset;
+            kernel (str) - 'linear' or 'rbf' kernel for computing the gram matrix;
+            unbiased (bool) - whether to compute an unbiased version of CKA;
+            sigma (float) - for 'rbf' kernel sigma defines the width of the Gaussian;
+        """
+        super().__init__(m=m, kernel=kernel, unbiased=unbiased, sigma=sigma)
+
+    def centering(self, K: Array) -> Array:
+        # The below code block is mainly copied from Simon Kornblith's implementation;
+        # see here: https://github.com/google-research/google-research/blob/master/representation_similarity/Demo.ipynb
+        """Centering of the (square) gram matrix K."""
+        if not np.allclose(K, K.T, rtol=1e-03, atol=1e-04):
+            raise ValueError("\nInput array must be a symmetric matrix.\n")
+        if self.unbiased:
+            # This formulation of the U-statistic, from Szekely, G. J., & Rizzo, M.
+            # L. (2014). Partial distance correlation with methods for dissimilarities.
+            # The Annals of Statistics, 42(6), 2382-2412, seems to be more numerically
+            # stable than the alternative from Song et al. (2007).
+            n = K.shape[0]
+            np.fill_diagonal(K, 0.0)
+            means = np.sum(K, axis=0, dtype=np.float64) / (n - 2)
+            means -= np.sum(means) / (2 * (n - 1))
+            K -= means[:, None]
+            K -= means[None, :]
+            np.fill_diagonal(K, 0.0)
+        else:
+            means = np.mean(K, axis=0, dtype=np.float64)
+            means -= np.mean(means) / 2
+            K -= means[:, None]
+            K -= means[None, :]
+        return K
+
+    def apply_kernel(self, X: Array) -> Array:
+        """Compute the (square) gram matrix K."""
+        try:
+            K = getattr(self, f"{self.kernel}_kernel")(X)
+        except AttributeError:
+            raise NotImplementedError
+        return K
+
+    def linear_kernel(self, X: Array) -> Array:
+        """Use a linear kernel for computing the gram matrix."""
+        return X @ X.T
+
+    def rbf_kernel(self, X: Array, sigma: float = 1.0) -> Array:
+        """Use an rbf kernel for computing the gram matrix. Sigma defines the width."""
+        GX = X @ X.T
+        KX = np.diag(GX) - GX + (np.diag(GX) - GX).T
+        if sigma is None:
+            mdist = np.median(KX[KX != 0])
+            sigma = math.sqrt(mdist)
+        KX *= -0.5 / sigma**2
+        KX = np.exp(KX)
+        return KX
+
+    def _hsic(self, X: Array, Y: Array) -> Array:
+        """Compute the Hilbert-Schmidt independence criterion."""
+        K = self.apply_kernel(X)
+        L = self.apply_kernel(Y)
+        K_c = self.centering(K)
+        L_c = self.centering(L)
+        # np.sum(K_c * L_c) is equivalent to K_c.flatten() @ L_c.flatten() or in math
+        # sum_{i=0}^{m} sum_{j=0}^{m} K^{\prime}_{ij} * L^{\prime}_{ij} = vec(K_c)^{T}vec(L_c)
+        return np.sum(K_c * L_c)
+
+    def compare(self, X: Array, Y: Array) -> Array:
+        """Compare two representation spaces X and Y using CKA."""
+        hsic_xy = self._hsic(X, Y)
+        hsic_xx = self._hsic(X, X)
+        hsic_yy = self._hsic(Y, Y)
+        rho = hsic_xy / np.sqrt(hsic_xx * hsic_yy)
+        return rho
diff --git a/thingsvision/core/cka/cka_torch.py b/thingsvision/core/cka/cka_torch.py
new file mode 100644
index 00000000..0144eede
--- /dev/null
+++ b/thingsvision/core/cka/cka_torch.py
@@ -0,0 +1,141 @@
+import re
+import warnings
+from typing import Optional, Union
+
+import torch
+from torchtyping import TensorType
+
+from .cka_base import CKABase
+
+
+class CKATorch(CKABase):
+    def __init__(
+        self,
+        m: int,
+        kernel: str,
+        unbiased: bool = False,
+        device: str = "cpu",
+        sigma: Optional[float] = 1.0,
+    ) -> None:
+        """
+        PyTorch implementation of CKA. Runs on CPU and CUDA.
+        Args:
+            m (int) - number of images / examples in a mini-batch or the full dataset;
+            kernel (str) - 'linear' or 'rbf' kernel for computing the gram matrix;
+            unbiased (bool) - whether to compute an unbiased version of CKA;
+            device (str) - whether to perform CKA computation on CPU or GPU;
+            sigma (float) - for 'rbf' kernel sigma defines the width of the Gaussian;
+        """
+        super().__init__(m=m, kernel=kernel, unbiased=unbiased, sigma=sigma)
+        device = self._check_device(device)
+        if device == "cpu":
+            self.hsic = self._hsic
+        else:
+            # use JIT compilation on CUDA
+            self.hsic = torch.compile(self._hsic)
+        self.device = torch.device(device)
+
+    @staticmethod
+    def _check_device(device: str) -> str:
+        """Check whether the selected device is available on current compute node."""
+        if device.startswith("cuda"):
+            gpu_index = re.search(r"cuda:(\d+)", device)
+
+            if not torch.cuda.is_available():
+                warnings.warn(
+                    "\nCUDA is not available on your system. Switching to device='cpu'.\n",
+                    category=UserWarning,
+                )
+                device = "cpu"
+            elif gpu_index and int(gpu_index.group(1)) >= torch.cuda.device_count():
+                warnings.warn(
+                    f"\nGPU index {gpu_index.group(1)} is out of range. "
+                    f"Available GPUs: {torch.cuda.device_count()}. "
+                    f"Switching to device='cuda:0'.\n",
+                    category=UserWarning,
+                )
+                device = "cuda:0"
+
+        print(f"\nUsing device: {device}\n")
+        return device
+
+    def centering(self, K: TensorType["m", "m"]) -> TensorType["m", "m"]:
+        """Centering of the (square) gram matrix K."""
+        # The below code block is mainly copied from Simon Kornblith's implementation;
+        # see here: https://github.com/google-research/google-research/blob/master/representation_similarity/Demo.ipynb
+        if not torch.allclose(K, K.T, rtol=1e-03, atol=1e-04):
+            raise ValueError("\nInput array must be a symmetric matrix.\n")
+        if self.unbiased:
+            # This formulation of the U-statistic, from Szekely, G. J., & Rizzo, M.
+            # L. (2014). Partial distance correlation with methods for dissimilarities.
+            # The Annals of Statistics, 42(6), 2382-2412, seems to be more numerically
+            # stable than the alternative from Song et al. (2007).
+            n = K.shape[0]
+            K.fill_diagonal_(0.0)
+            means = K.sum(dim=0) / (n - 2)
+            means -= means.sum() / (2 * (n - 1))
+            K -= means[:, None]
+            K -= means[None, :]
+            K.fill_diagonal_(0.0)
+        else:
+            means = K.mean(dim=0)
+            means -= means.mean() / 2
+            K -= means[:, None]
+            K -= means[None, :]
+        return K
+
+    def apply_kernel(
+        self, X: Union[TensorType["m", "d"], TensorType["m", "p"]]
+    ) -> TensorType["m", "m"]:
+        """Compute the (square) gram matrix K."""
+        try:
+            K = getattr(self, f"{self.kernel}_kernel")(X)
+        except AttributeError:
+            raise NotImplementedError
+        return K
+
+    def linear_kernel(
+        self, X: Union[TensorType["m", "d"], TensorType["m", "p"]]
+    ) -> TensorType["m", "m"]:
+        """Use a linear kernel for computing the gram matrix."""
+        return X @ X.T
+
+    def rbf_kernel(
+        self, X: TensorType["m", "d"], sigma: Optional[float] = 1.0
+    ) -> TensorType["m", "m"]:
+        """Use an rbf kernel for computing the gram matrix. Sigma defines the width."""
+        GX = X @ X.T
+        KX = torch.diag(GX) - GX + (torch.diag(GX) - GX).T
+        if sigma is None:
+            mdist = torch.median(KX[KX != 0])
+            sigma = torch.sqrt(mdist)
+        KX *= -0.5 / sigma**2
+        KX = KX.exp()
+        return KX
+
+    def _hsic(
+        self, X: TensorType["m", "d"], Y: TensorType["m", "p"]
+    ) -> TensorType["1"]:
+        K = self.apply_kernel(X)
+        L = self.apply_kernel(Y)
+        K_c = self.centering(K)
+        L_c = self.centering(L)
+        """Compute the Hilbert-Schmidt independence criterion."""
+        # np.sum(K_c * L_c) is equivalent to K_c.flatten() @ L_c.flatten() or in math
+        # sum_{i=0}^{m} sum_{j=0}^{m} K^{\prime}_{ij} * L^{\prime}_{ij} = vec(K_c)^{T}vec(L_c)
+        return torch.sum(K_c * L_c)
+
+    @torch.inference_mode()
+    def compare(
+        self, X: TensorType["m", "d"], Y: TensorType["m", "p"]
+    ) -> TensorType["1"]:
+        """Compare two representation spaces X and Y using CKA."""
+        X = X.to(self.device)
+        Y = Y.to(self.device)
+        hsic_xy = self.hsic(X, Y)
+        hsic_xx = self.hsic(X, X)
+        hsic_yy = self.hsic(Y, Y)
+        rho = hsic_xy / torch.sqrt(hsic_xx * hsic_yy)
+        if rho.is_cuda:
+            rho = rho.cpu()
+        return rho
diff --git a/thingsvision/core/cka/helpers.py b/thingsvision/core/cka/helpers.py
new file mode 100644
index 00000000..010b6063
--- /dev/null
+++ b/thingsvision/core/cka/helpers.py
@@ -0,0 +1,28 @@
+from typing import Optional, Union
+
+from .cka_numpy import CKANumPy
+from .cka_torch import CKATorch
+
+BACKENDS = ["numpy", "torch"]
+
+
+def get_cka(
+    backend: str,
+    m: int,
+    kernel: str = "linear",
+    unbiased: bool = False,
+    sigma: Optional[float] = 1.0,
+    device: Optional[str] = None,
+) -> Union[CKANumPy, CKATorch]:
+    """Return a NumPy or PyTorch implementation of CKA."""
+    assert backend in BACKENDS, f"\nSupported backends are: {BACKENDS}\n"
+    if backend == "numpy":
+        cka = CKANumPy(m=m, kernel=kernel, unbiased=unbiased, sigma=sigma)
+    else:
+        assert isinstance(
+            device, str
+        ), "\nDevice must be set for using PyTorch backend.\n"
+        cka = CKATorch(
+            m=m, kernel=kernel, unbiased=unbiased, device=device, sigma=sigma
+        )
+    return cka


Test patch:
diff --git a/tests/helper.py b/tests/helper.py
index 7d3a6dc2..d8dd2243 100644
--- a/tests/helper.py
+++ b/tests/helper.py
@@ -9,7 +9,6 @@
 import torch.nn as nn
 from tensorflow.keras.layers import Dense
 from tensorflow.keras.models import Sequential
-
 from thingsvision import get_extractor
 from thingsvision.utils.data import DataLoader, ImageDataset
 
diff --git a/tests/test_features.py b/tests/test_features.py
index 0d48f87e..d97d1f48 100644
--- a/tests/test_features.py
+++ b/tests/test_features.py
@@ -1,9 +1,9 @@
 import os
 import shutil
 import unittest
-import torch
 
 import numpy as np
+import torch
 from thingsvision.core.extraction import center_features, normalize_features
 from thingsvision.utils.storing import save_features, split_features
 
diff --git a/tests/test_rest.py b/tests/test_rest.py
index db6b40e4..2ddc01e2 100644
--- a/tests/test_rest.py
+++ b/tests/test_rest.py
@@ -2,12 +2,14 @@
 import unittest
 
 import numpy as np
+import torch
+from thingsvision.core.cka import get_cka
 from thingsvision.core.rsa import compute_rdm, correlate_rdms, plot_rdm
-from thingsvision.core.cka import CKA
 from thingsvision.utils.storing import save_features
 
 import tests.helper as helper
 
+
 class RSATestCase(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -48,7 +50,6 @@ def test_rdms(self):
             self.assertTrue(corr > float(-1) and corr < float(1))
 
 
-
 class CKATestCase(unittest.TestCase):
 
     @classmethod
@@ -75,11 +76,28 @@ def test_cka(self):
         )
         self.assertEqual(features_i.shape, features_j.shape)
         m = features_i.shape[0]
-        for kernel in ['linear', 'rbf']:
-            cka = CKA(m=m, kernel=kernel, sigma=0.5 if kernel == 'rbf' else None)
-            rho = cka.compare(features_i, features_j)
-            self.assertTrue(isinstance(rho, float))
-            self.assertTrue(rho > float(-1) and rho < float(1))
+        for backend in ["numpy", "torch"]:
+            if backend == "torch":
+                device = "cpu"
+                features_i = torch.from_numpy(features_i)
+                features_j = torch.from_numpy(features_j)
+            else:
+                device = None
+            for kernel in ["linear", "rbf"]:
+                sigma = 0.5 if kernel == "rbf" else None
+                for debiased in [True, False]:
+                    cka = get_cka(
+                        backend=backend,
+                        m=m,
+                        unbiased=debiased,
+                        kernel=kernel,
+                        sigma=sigma,
+                        device=device,
+                    )
+                    rho = cka.compare(features_i, features_j)
+                    if backend == "torch":
+                        rho = rho.item()
+                    self.assertTrue(rho > float(-1) and rho < float(1))
 
 
 class FileNamesTestCase(unittest.TestCase):