dataeval 0.84.0__py3-none-any.whl → 0.84.1__py3-none-any.whl

{dataeval-0.84.0.dist-info → dataeval-0.84.1.dist-info}/RECORD

@@ -1,20 +1,19 @@
-dataeval/__init__.py,sha256=VczdyekiNdqvi2dEUf7xXBu3Aw-MoVTvq-k6c2zjeBM,1636
+dataeval/__init__.py,sha256=QzrctVrymZuLN8tnHcF1wp0RTXYM3WSWMozX3NOzIos,1636
 dataeval/_log.py,sha256=Mn5bRWO0cgtAYd5VGYSFiPgu57ta3zoktrtHAZ1m3dU,357
 dataeval/config.py,sha256=lD1YDH8HosFeRU5rQEYRBcmXMZy-csWaMlJTRZGd9iU,3582
 dataeval/detectors/__init__.py,sha256=3Sg-XWlwr75zEEH3hZKA4nWMtGvaRlnfzTWvZG_Ak6U,189
-dataeval/detectors/drift/__init__.py,sha256=6is_XBtG1d-vUbhHvqXGOdnAwxJ7NA5yRfURn7pCeIw,651
-dataeval/detectors/drift/_base.py,sha256=mJdKvyROgWvz-p1VlAIJqUI6BAj9ss8riUvR5An5wIw,13459
-dataeval/detectors/drift/_cvm.py,sha256=H2w-I0eMD7yP-CSmpdodeJ0-TYznJT7w_H7JuobESow,3859
-dataeval/detectors/drift/_ks.py,sha256=-5k3RBPA3kadX7oD14Wc52rAqQf1udwFeW7Qf3Sv4Tw,4058
-dataeval/detectors/drift/_mmd.py,sha256=NEXowx9UHIvmEKS8sqssw6PMLJMh0BZPhRNX1hYlkz4,7239
-dataeval/detectors/drift/_torch.py,sha256=VrFCyTaRrUslFPy_mYZ4UL70LZ8faH4eHwLurZ9qqNE,7628
-dataeval/detectors/drift/_uncertainty.py,sha256=O5h6_bJbeQEE660SLLP8k-EHqImmKegIcxzcnUKI7X4,5714
-dataeval/detectors/drift/updates.py,sha256=Btu2iaZW7fbO59G1w5v3ykFot0YPzy2U6VjF0d440VE,2195
+dataeval/detectors/drift/__init__.py,sha256=gD8aY5PotS-S2ot7iB_z_zzSOjIbQLw5znFBNj0jtHE,646
+dataeval/detectors/drift/_base.py,sha256=PdWyEuYqExFdyxvyOh7Q8yXnjNm0D3KfpDUn0bUixtY,7580
+dataeval/detectors/drift/_cvm.py,sha256=CSEyNN9u1MzUI6QmCSlexTUSlHzK1kYh36Nv2L72WbY,3016
+dataeval/detectors/drift/_ks.py,sha256=ifFb_0JcyykJyF9DAVkQqWCXc-3aA0AC8c8to_oOPKo,3198
+dataeval/detectors/drift/_mmd.py,sha256=DOHBNyNNxosR67yM9HTxbvqp1IZ8_KSvTVlX-JtKtjM,11601
+dataeval/detectors/drift/_uncertainty.py,sha256=BHlykJ-r7TGLJxdPfoazXnoAJ1qVDzbk5HjAMdsnHz8,5847
+dataeval/detectors/drift/updates.py,sha256=L1PnrPlIE1x6ujCc5mCwjcAZwadVTn-Zjb6MnTDvzJQ,2251
 dataeval/detectors/linters/__init__.py,sha256=xn2zPwUcmsuf-Jd9uw6AVI11C9z1b1Y9fYtuFnXenZ0,404
-dataeval/detectors/linters/duplicates.py,sha256=tcxniL8rRZkDdQqfuS502UmfKxS3a7iRA22Dtt_vQIk,4935
-dataeval/detectors/linters/outliers.py,sha256=Hln2dPQZjF_uV2QYptA_o6ZF3ugyCImVT-XLDB2-q3A,9042
+dataeval/detectors/linters/duplicates.py,sha256=x36-0EAlO_AuOttvElJOZCa0R3VzrlII0NxjwhdkrpE,4969
+dataeval/detectors/linters/outliers.py,sha256=Z0Sbtluu2im1IRGsjhXF2AhrShKDrEkF8BWzAZyPwlA,9054
 dataeval/detectors/ood/__init__.py,sha256=juCYBDs7CQEAtMhnEpPqF6uTrOIH9kTBSuQ_GRw6a8o,283
-dataeval/detectors/ood/ae.py,sha256=YQfhB1ShQLjM1V4uCz9Oo2tCZpOfAZ_-SBCAl4Ac67Y,2921
+dataeval/detectors/ood/ae.py,sha256=fTrUfFxv6xUqzKpwMC8rW3JrizA16M_bgzqLuBKMrS0,2944
 dataeval/detectors/ood/base.py,sha256=9b-Ljznf0lB1SXF4F_Aj3eJ4Y3ijGEDPMjucUsWOGJM,3051
 dataeval/detectors/ood/mixin.py,sha256=0_o-1HPvgf3-Lf1MSOIfjj5UB8LTLEBGYtJJfyCCzwc,5431
 dataeval/detectors/ood/vae.py,sha256=Fcq0-WbLhzYCgYOAJPBklHm7yuXmFJuEpBkhgwM5kiA,2291
@@ -35,14 +34,14 @@ dataeval/metrics/estimators/_clusterer.py,sha256=1HrpihGTJ63IkNSOy4Ibw633Gllkm1R
 dataeval/metrics/estimators/_divergence.py,sha256=QDWl1lyAYoO9D3Ho7qOHSk6ud8Gi2MGuXEsYwO1HxvA,4043
 dataeval/metrics/estimators/_uap.py,sha256=BULEBbJ9BQ1IcTeZf0x7iI60QHAWCccBOM97FIu9VXA,1928
 dataeval/metrics/stats/__init__.py,sha256=6tA_9nbbM5ObJ6cds8Y1VBtTQiTOxrpGQSFLu_lWGGA,1098
-dataeval/metrics/stats/_base.py,sha256=rA-Xt9slf2DOR5ky9gGR5s1pmzTb47DykovDp5EWEP0,10672
+dataeval/metrics/stats/_base.py,sha256=YIfOVGd7E19B4dpAnzDYRQkaikvRRyJIpznJNfVtPdw,10750
 dataeval/metrics/stats/_boxratiostats.py,sha256=8Kd2FTZ5PLNYZfdAjU_R385gb0Z16JY0L9H_d5ZhgQs,6341
-dataeval/metrics/stats/_dimensionstats.py,sha256=h2wCLn4UuW7-GV6tM5E1SqSeGa_-4ie9oaEXpSC7EKI,2690
-dataeval/metrics/stats/_hashstats.py,sha256=yD6cXMvOo10-xtwUr7ftBRbCqMhReNfQJMInEWV_8Mk,4757
-dataeval/metrics/stats/_imagestats.py,sha256=hyjijPXAfUIJ1lwWiIyYK9VSLiq7Vg2-YhJ5Q8s1rkY,2979
+dataeval/metrics/stats/_dimensionstats.py,sha256=73mFP-Myxne0peFliwvTntc0kk4cpq0krzMvSLDSIMM,2702
+dataeval/metrics/stats/_hashstats.py,sha256=gp9X_pnTT3mPH9YNrWLdn2LQPK_epJ3dQRoyOCwmKlg,4758
+dataeval/metrics/stats/_imagestats.py,sha256=gUPNgN5Zwzdr7WnSwbve1NXNsyxd5dy3cSnlR_7guCg,3007
 dataeval/metrics/stats/_labelstats.py,sha256=WbvXZ831a5BDfm58HF8Z8i5JUV1tgw7tcfzUh8pOXSo,2825
-dataeval/metrics/stats/_pixelstats.py,sha256=Q0-ldG-znDYBP_qTqm6S4qYm0ZV5FTTHf8MlyGHSYEc,3235
-dataeval/metrics/stats/_visualstats.py,sha256=ZxBDTerZ8ixibY2pGl7mwwcIz3DWl-k_Jb4YwBjHLNw,3686
+dataeval/metrics/stats/_pixelstats.py,sha256=SfergRbjNJE4h0xqe-0c8RnKtZmEkZ9MwExdipLSGvg,3247
+dataeval/metrics/stats/_visualstats.py,sha256=cq4AbF2B50Ihbzb86FphcnKQ1TSwNnP3PsnbpiPQZWw,3698
 dataeval/outputs/__init__.py,sha256=ciK-RdXgtn_s7MSCUW1UXvrXltMbltqbpfe9_V7xGrI,1701
 dataeval/outputs/_base.py,sha256=aZFbgybnZSQ3ws7QYRLTbDFqUfBFRVtIwX2LZfeGFUA,5703
 dataeval/outputs/_bias.py,sha256=GwbjLdppUODOeudYb_7ki2ejDmAYthlRKGijVwgVePE,12407
@@ -55,9 +54,9 @@ dataeval/outputs/_stats.py,sha256=c73Yc3Kkrl-MN6BGKe1V0Yr6Ix2Yp_DZZfFSp8fZMZ0,13
 dataeval/outputs/_utils.py,sha256=HHlGC7sk416m_3Bgn075Qdblz_aPup_UOafJpB0RuXY,893
 dataeval/outputs/_workflows.py,sha256=MkRD6ubI4NCBXb9v3kjXy64cUGs3G-JKkBdOpRD9XVE,10750
 dataeval/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dataeval/typing.py,sha256=h1wNoWovasIEAWwNWvICToJ1bE1DRba68HifP14zjZc,6598
+dataeval/typing.py,sha256=zn6smomSdcO7EeZpeeSP5-8sknTdgUuU7TKe-3UFVrg,6550
 dataeval/utils/__init__.py,sha256=T8F8zJh4ZAeu0wDzfpld92I2zJg9mWBmkGCHrDPU7gk,264
-dataeval/utils/_array.py,sha256=B44LbVi4g7XntVvnPyi5KeyQIxRQSvYhTVD4Th76u5g,5577
+dataeval/utils/_array.py,sha256=KqAdXEMjcXYvdWdYEEoEbigwQJ4S9VYxQS3sRFeY5XY,5929
 dataeval/utils/_bin.py,sha256=nylthmsC3vzLHLhlUMACvZs--h7xvAh9Pt75InaQJW8,7322
 dataeval/utils/_clusterer.py,sha256=fw5x-2QN0TIbiodDKHZxRgxKHINedpPcOklzce0Rbjg,5436
 dataeval/utils/_fast_mst.py,sha256=4_7ykVihCL5jWtxcGnrecIsDQo65kUml9SZ1JxgBZYY,7172
@@ -66,41 +65,42 @@ dataeval/utils/_method.py,sha256=9B9JQbgqWJBRhQJb7glajUtWaQzUTIUuvrZ9_bisxsM,394
 dataeval/utils/_mst.py,sha256=f0vXytTUjlOS6AyL7c6PkXmaHuuGUK-vMLpq-5xMgxk,2183
 dataeval/utils/_plot.py,sha256=mTRQNbJsA42QMiOwZbJaH8sNYgP996QFDEGVVE9HSgY,7076
 dataeval/utils/data/__init__.py,sha256=vldQ2ZXl8gnI3s4vAGqUUVi6dc_R58F3JMSpbCOyFRI,820
-dataeval/utils/data/_dataset.py,sha256=RZT05cfXkiPJPNCG6SVf8zmsk0pQUBViyrSo2l1_G5w,7852
-dataeval/utils/data/_embeddings.py,sha256=NK87PfzpQUagwU1aBknsEEihAPNR3BIqHnHkpeKEgVs,7028
-dataeval/utils/data/_images.py,sha256=Zn2um-oZjypwVTdpQNw7DsjxKyujswe2jLIgxmUPQ7Q,2626
-dataeval/utils/data/_metadata.py,sha256=VqeePp7NtoFFWzmIhH4fn-cjrnATpgzgzs-d73cnBXM,14370
+dataeval/utils/data/_dataset.py,sha256=MHY582yRm4FxQkkLWUhKZBb7ZyvWypM6ldUG89vd3uE,7936
+dataeval/utils/data/_embeddings.py,sha256=iDtfLJY1uHoTP4UdQoOt-3wopc6kSOXH_4CVNnmXXA4,8356
+dataeval/utils/data/_images.py,sha256=WF9XJRka8ohUdyI2IKBMAy3JoJhOm1iC-8tbYl8woRM,2642
+dataeval/utils/data/_metadata.py,sha256=62z_qHjoGjiMdpuT36QpNhbWy2UClHWUcjHHlIWT470,14464
 dataeval/utils/data/_selection.py,sha256=2c6DjyeDIJapbI7xL36eBxFnJHIP8Yxt3oU3rBGMqLk,3948
 dataeval/utils/data/_split.py,sha256=q-2RwllJgazwuyxB_GoBqK_nLkqIjyTVr2SQKj_7lhw,16767
 dataeval/utils/data/_targets.py,sha256=ws5d9wRiDkIuOV7GSAKNxzgSm6AWTgb0BFroQK5nAmM,3057
 dataeval/utils/data/collate.py,sha256=Z5nmBnWV_IoJzMp_tj8RCKjMJA9sSCY_zZITqISGixc,3865
 dataeval/utils/data/datasets/__init__.py,sha256=jBrswiERrvBx4pJQJZIq_B5UE-Wy8a2_SBfM2crG8R8,511
-dataeval/utils/data/datasets/_base.py,sha256=CZ-hb-yWPLdnTQ3pURJMcityQ42ZNYj_Lbb1P5Junn4,8793
-dataeval/utils/data/datasets/_cifar10.py,sha256=I6HKksE2escos1aTdiZJObtiVXChBlez5BDa0eBfJ_Y,5449
+dataeval/utils/data/datasets/_base.py,sha256=827nSVhZ-tqeHw1HQ7Qj060CSDd90fEWZomN6FaWnQA,8872
+dataeval/utils/data/datasets/_cifar10.py,sha256=R7QgcCHowAkqhEXOvUhybXTmMlA4BJXkTuAeV9uDgfU,5449
 dataeval/utils/data/datasets/_fileio.py,sha256=SixIk5nIlIwJdX9zjNXS10vHA3hL8aaYbqHsDg1xSpY,6447
-dataeval/utils/data/datasets/_milco.py,sha256=ScBe7Ux-J9Kxs33jeKffhWKeSb8GCrWznTyEUt95Vt4,6369
+dataeval/utils/data/datasets/_milco.py,sha256=bVVDl5W8TdTPU2RiwoPXrfFDM1rKyb-LslwTThBXEr0,7583
 dataeval/utils/data/datasets/_mixin.py,sha256=FJgZP_cpJkgAHA3j3ai_j3Wt7aFSEjIMVmt9NpvVXzg,1757
-dataeval/utils/data/datasets/_mnist.py,sha256=iWWI9mq6TbZm7eTL9btzqjCNMhgXrLHQeMKENr7USsk,7988
-dataeval/utils/data/datasets/_ships.py,sha256=p3fScYLW2f1wUEPOroCX5nOFti0vMOSjeYltj6ox53U,4777
+dataeval/utils/data/datasets/_mnist.py,sha256=kNDJw0oyqa6QgU1y9lg-3AzStavK1BB8iHnDOdv9nyE,8112
+dataeval/utils/data/datasets/_ships.py,sha256=rsyIoRAIk40liFgaEb2dg0lYB7__bAGd9zh9ouzjFKg,4880
 dataeval/utils/data/datasets/_types.py,sha256=iSKyHXRlGuomXs0FHK6md8lXLQrQQ4fxgVOwr4o81bo,1089
-dataeval/utils/data/datasets/_voc.py,sha256=4poEer_G_mUBcz6eAro0Tc29CjdgjEAlms0Eu0tLBzE,14842
-dataeval/utils/data/selections/__init__.py,sha256=k86OpqGPkjT1MrOir5fOZ3AIq5UR81Az9ek7l1-GdIM,565
-dataeval/utils/data/selections/_classfilter.py,sha256=opSF8CGv4x1hUMe-GTQOu3UwJK80DzT0nJOV0l2uaW4,2404
+dataeval/utils/data/datasets/_voc.py,sha256=QUtpbh2EpiBoicsmOo-YIfwRwPXyHj-zB2hFn7tlz0Y,15580
+dataeval/utils/data/selections/__init__.py,sha256=iUbMZRDuBXwY3SNAtZTdCVu7SI4zbCyaL6ItXnnq1yI,655
+dataeval/utils/data/selections/_classbalance.py,sha256=hHq9frdwzFLCUmfeJq977Sot_SXhuGANlSsetokhRDc,1465
+dataeval/utils/data/selections/_classfilter.py,sha256=xdR5uX7W5Yivf-mE_CikbRi2fGrZLFrPYun3TeQHTA0,1267
 dataeval/utils/data/selections/_indices.py,sha256=QdLgXN7GABCvGPYe28PV1RAc_RSP_nZOyCvEpKRBdWg,636
 dataeval/utils/data/selections/_limit.py,sha256=ECvHRsp7OF4LZw2tE4sGqqJ085kjC-hd2c7QDMfvXr8,518
-dataeval/utils/data/selections/_prioritize.py,sha256=EAA4_uFVV7MmemhhufGmP7eunnbtyTc-TzgcnvRK5OE,11333
+dataeval/utils/data/selections/_prioritize.py,sha256=uRQjeQiAc-vvwHMH4CQtXTGJCfjj_h5dgGlhQYFMz1c,11318
 dataeval/utils/data/selections/_reverse.py,sha256=6SWpELC9Wgx-kPqzhDrPNn4NKU6FqDJveLrxV4D2Ypk,374
-dataeval/utils/data/selections/_shuffle.py,sha256=kY3xJvVbBArdrJu_u6mXmxk1HdNmmDE4w7MmxbevUmU,1178
+dataeval/utils/data/selections/_shuffle.py,sha256=_jwms0qcwrknf2Fx84cCXyNOJyhE_V8rcnDOTDn1S2A,1179
 dataeval/utils/metadata.py,sha256=1XeGYj_e97-nJ_IrWEHPhWICmouYU5qbXWbp7uhZrIE,14171
 dataeval/utils/torch/__init__.py,sha256=dn5mjCrFp0b1aL_UEURhONU0Ag0cmXoTOBSGagpkTiA,325
 dataeval/utils/torch/_blocks.py,sha256=HVhBTMMD5NA4qheMUgyol1KWiKZDIuc8k5j4RcMKmhk,1466
 dataeval/utils/torch/_gmm.py,sha256=XM68GNEP97EjaB1U49-ZXRb81d0CEFnPS910alrcB3g,3740
-dataeval/utils/torch/_internal.py,sha256=23DCnF7C7N3tZgZUpT2nyH7mMb8Pi4GcnQyjK0BKHpg,5735
+dataeval/utils/torch/_internal.py,sha256=vHy-DzPhmvE8h3wmWc3aciBJ8nDGzQ1z1jTZgGjmDyM,4154
 dataeval/utils/torch/models.py,sha256=hmroEs6C6jQ5tAoZa71RFeIvXLxfXrTJSFH_jG2LGQU,9749
 dataeval/utils/torch/trainer.py,sha256=iUotX4OdirH8-ZtjdpU8gbJavkYW9YY9qpA2mAlFy1Y,5520
 dataeval/workflows/__init__.py,sha256=ou8y0KO-d6W5lgmcyLjKlf-J_ckP3vilW7wHkgiDlZ4,255
 dataeval/workflows/sufficiency.py,sha256=mjKmfRrAjShLUFIARv5o8yT5fnFvDsS5Qu6ujIPUgQg,8497
-dataeval-0.84.0.dist-info/LICENSE.txt,sha256=uAooygKWvX6NbU9Ran9oG2msttoG8aeTeHSTe5JeCnY,1061
-dataeval-0.84.0.dist-info/METADATA,sha256=oEAANNRg8RUIWn9AdrQEV7OUnX5mJbgf4NqXr5QY8AY,5320
-dataeval-0.84.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-dataeval-0.84.0.dist-info/RECORD,,
+dataeval-0.84.1.dist-info/LICENSE.txt,sha256=uAooygKWvX6NbU9Ran9oG2msttoG8aeTeHSTe5JeCnY,1061
+dataeval-0.84.1.dist-info/METADATA,sha256=F7L5PSWHV3z0_4pwA-JSgucW2A4bEv_dtvIMzCTGLZ8,5308
+dataeval-0.84.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+dataeval-0.84.1.dist-info/RECORD,,

dataeval/detectors/drift/_torch.py

@@ -1,222 +0,0 @@
-"""
-Source code derived from Alibi-Detect 0.11.4
-https://github.com/SeldonIO/alibi-detect/tree/v0.11.4
-
-Original code Copyright (c) 2023 Seldon Technologies Ltd
-Licensed under Apache Software License (Apache 2.0)
-"""
-
-from __future__ import annotations
-
-__all__ = []
-
-from typing import Any, Callable
-
-import numpy as np
-import torch
-import torch.nn as nn
-from numpy.typing import NDArray
-
-from dataeval.config import DeviceLike, get_device
-from dataeval.utils.torch._internal import predict_batch
-
-
-def mmd2_from_kernel_matrix(
-    kernel_mat: torch.Tensor, m: int, permute: bool = False, zero_diag: bool = True
-) -> torch.Tensor:
-    """
-    Compute maximum mean discrepancy (MMD^2) between 2 samples x and y from the
-    full kernel matrix between the samples.
-
-    Parameters
-    ----------
-    kernel_mat : torch.Tensor
-        Kernel matrix between samples x and y.
-    m : int
-        Number of instances in y.
-    permute : bool, default False
-        Whether to permute the row indices. Used for permutation tests.
-    zero_diag : bool, default True
-        Whether to zero out the diagonal of the kernel matrix.
-
-    Returns
-    -------
-    torch.Tensor
-        MMD^2 between the samples from the kernel matrix.
-    """
-    n = kernel_mat.shape[0] - m
-    if zero_diag:
-        kernel_mat = kernel_mat - torch.diag(kernel_mat.diag())
-    if permute:
-        idx = torch.randperm(kernel_mat.shape[0])
-        kernel_mat = kernel_mat[idx][:, idx]
-    k_xx, k_yy, k_xy = kernel_mat[:-m, :-m], kernel_mat[-m:, -m:], kernel_mat[-m:, :-m]
-    c_xx, c_yy = 1 / (n * (n - 1)), 1 / (m * (m - 1))
-    mmd2 = c_xx * k_xx.sum() + c_yy * k_yy.sum() - 2.0 * k_xy.mean()
-    return mmd2
-
-
-def preprocess_drift(
-    x: NDArray[Any],
-    model: nn.Module,
-    device: DeviceLike | None = None,
-    preprocess_batch_fn: Callable | None = None,
-    batch_size: int = int(1e10),
-    dtype: type[np.generic] | torch.dtype = np.float32,
-) -> NDArray[Any] | torch.Tensor | tuple[Any, ...]:
-    """
-    Prediction function used for preprocessing step of drift detector.
-
-    Parameters
-    ----------
-    x : NDArray
-        Batch of instances.
-    model : nn.Module
-        Model used for preprocessing.
-    device : DeviceLike or None, default None
-        The hardware device to use if specified, otherwise uses the DataEval
-        default or torch default.
-    preprocess_batch_fn : Callable or None, default None
-        Optional batch preprocessing function. For example to convert a list of objects
-        to a batch which can be processed by the PyTorch model.
-    batch_size : int, default 1e10
-        Batch size used during prediction.
-    dtype : np.dtype or torch.dtype, default np.float32
-        Model output type, either a :term:`NumPy` or torch dtype, e.g. np.float32 or torch.float32.
-
-    Returns
-    -------
-    NDArray | torch.Tensor | tuple
-        Numpy array, torch tensor or tuples of those with model outputs.
-    """
-    return predict_batch(
-        x,
-        model,
-        device=get_device(device),
-        batch_size=batch_size,
-        preprocess_fn=preprocess_batch_fn,
-        dtype=dtype,
-    )
-
-
-@torch.jit.script
-def _squared_pairwise_distance(
-    x: torch.Tensor, y: torch.Tensor, a_min: float = 1e-30
-) -> torch.Tensor:  # pragma: no cover - torch.jit.script code is compiled and copied
-    """
-    PyTorch pairwise squared Euclidean distance between samples x and y.
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        Batch of instances of shape [Nx, features].
-    y : torch.Tensor
-        Batch of instances of shape [Ny, features].
-    a_min : float
-        Lower bound to clip distance values.
-
-    Returns
-    -------
-    torch.Tensor
-        Pairwise squared Euclidean distance [Nx, Ny].
-    """
-    x2 = x.pow(2).sum(dim=-1, keepdim=True)
-    y2 = y.pow(2).sum(dim=-1, keepdim=True)
-    dist = torch.addmm(y2.transpose(-2, -1), x, y.transpose(-2, -1), alpha=-2).add_(x2)
-    return dist.clamp_min_(a_min)
-
-
-def sigma_median(x: torch.Tensor, y: torch.Tensor, dist: torch.Tensor) -> torch.Tensor:
-    """
-    Bandwidth estimation using the median heuristic `Gretton2012`
-
-    Parameters
-    ----------
-    x : torch.Tensor
-        Tensor of instances with dimension [Nx, features].
-    y : torch.Tensor
-        Tensor of instances with dimension [Ny, features].
-    dist : torch.Tensor
-        Tensor with dimensions [Nx, Ny], containing the pairwise distances
-        between `x` and `y`.
-
-    Returns
-    -------
-    torch.Tensor
-        The computed bandwidth, `sigma`.
-    """
-    n = min(x.shape[0], y.shape[0])
-    n = n if (x[:n] == y[:n]).all() and x.shape == y.shape else 0
-    n_median = n + (np.prod(dist.shape) - n) // 2 - 1
-    sigma = (0.5 * dist.flatten().sort().values[int(n_median)].unsqueeze(dim=-1)) ** 0.5
-    return sigma
-
-
-class GaussianRBF(nn.Module):
-    """
-    Gaussian RBF kernel: k(x,y) = exp(-(1/(2*sigma^2)||x-y||^2).
-
-    A forward pass takes a batch of instances x [Nx, features] and
-    y [Ny, features] and returns the kernel matrix [Nx, Ny].
-
-    Parameters
-    ----------
-    sigma : torch.Tensor | None, default None
-        Bandwidth used for the kernel. Needn't be specified if being inferred or
-        trained. Can pass multiple values to eval kernel with and then average.
-    init_sigma_fn : Callable | None, default None
-        Function used to compute the bandwidth ``sigma``. Used when ``sigma`` is to be
-        inferred. The function's signature should take in the tensors ``x``, ``y`` and
-        ``dist`` and return ``sigma``. If ``None``, it is set to ``sigma_median``.
-    trainable : bool, default False
-        Whether or not to track gradients w.r.t. `sigma` to allow it to be trained.
-    """
-
-    def __init__(
-        self,
-        sigma: torch.Tensor | None = None,
-        init_sigma_fn: Callable | None = None,
-        trainable: bool = False,
-    ) -> None:
-        super().__init__()
-        init_sigma_fn = sigma_median if init_sigma_fn is None else init_sigma_fn
-        self.config: dict[str, Any] = {
-            "sigma": sigma,
-            "trainable": trainable,
-            "init_sigma_fn": init_sigma_fn,
-        }
-        if sigma is None:
-            self.log_sigma: nn.Parameter = nn.Parameter(torch.empty(1), requires_grad=trainable)
-            self.init_required: bool = True
-        else:
-            sigma = sigma.reshape(-1)  # [Ns,]
-            self.log_sigma: nn.Parameter = nn.Parameter(sigma.log(), requires_grad=trainable)
-            self.init_required: bool = False
-        self.init_sigma_fn = init_sigma_fn
-        self.trainable = trainable
-
-    @property
-    def sigma(self) -> torch.Tensor:
-        return self.log_sigma.exp()
-
-    def forward(
-        self,
-        x: np.ndarray[Any, Any] | torch.Tensor,
-        y: np.ndarray[Any, Any] | torch.Tensor,
-        infer_sigma: bool = False,
-    ) -> torch.Tensor:
-        x, y = torch.as_tensor(x), torch.as_tensor(y)
-        dist = _squared_pairwise_distance(x.flatten(1), y.flatten(1))  # [Nx, Ny]
-
-        if infer_sigma or self.init_required:
-            if self.trainable and infer_sigma:
-                raise ValueError("Gradients cannot be computed w.r.t. an inferred sigma value")
-            sigma = self.init_sigma_fn(x, y, dist)
-            with torch.no_grad():
-                self.log_sigma.copy_(sigma.log().clone())
-            self.init_required: bool = False
-
-        gamma = 1.0 / (2.0 * self.sigma**2)  # [Ns,]
-        # TODO: do matrix multiplication after all?
-        kernel_mat = torch.exp(-torch.cat([(g * dist)[None, :, :] for g in gamma], dim=0))  # [Ns, Nx, Ny]
-        return kernel_mat.mean(dim=0)  # [Nx, Ny]