kerneltests 0.0.1__py3-none-any.whl

kerneltests-0.0.1.dist-info/METADATA

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: kerneltests
+Version: 0.0.1
+Summary: Unified, tested implementations of modern kernel hypothesis tests (MMDAgg, ...)
+Author: Ankith Burki
+License: MIT
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.22
+Requires-Dist: scipy>=1.8
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Provides-Extra: jax
+Requires-Dist: jax>=0.4; extra == "jax"
+Provides-Extra: docs
+Requires-Dist: mkdocs; extra == "docs"
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocstrings[python]; extra == "docs"
+Dynamic: license-file
+
+# kerneltests
+
+`kerneltests` (import name `kt`) gives you modern kernel hypothesis tests behind one small API: two-sample testing (**MMDAgg**, **MMDAggInc**, **MMD-FUSE**), goodness-of-fit testing (**KSDAgg**, **KSDAggInc**) and independence testing (**HSICAggInc**). Every port is checked against the original reference implementations by Schrab et al. and gives **bit-exact** agreement for the same seed (MMD-FUSE, whose reference is JAX, agrees to rtol 1e-6 given the same permutations) (identical decisions, statistics, quantiles, p-values and thresholds) when run with `fast=False`. The default `fast=True` path for equal sample sizes is 1.7–2.3× faster than the reference MMDAgg and agrees with it up to floating-point rounding (identical decisions in all tested cases; statistics within ~1e-12 relative). The ports also fix five confirmed bugs in the references (all listed in CHANGELOG.md), including:
+
+1. **MMDAgg p-values.** With `kernel="laplace_gaussian"`, the reference reports Gaussian-kernel p-values that duplicate the Laplace ones, because `mmdagg/np.py` line 352 indexes `p_vals[i]` where it should use `p_vals[idx]`.
+2. **KSDAgg internal assertion.** The reference can raise an `AssertionError` when its two equivalent rejection rules (`p-value <= threshold` and `KSD > quantile`) disagree because of floating-point rounding.
+
+## Install
+
+```bash
+pip install git+<repository-url>      # or, from a clone:
+pip install -e .
+```
+
+You need Python 3.9 or newer, NumPy 1.22 or newer, and SciPy 1.8 or newer.
+
+Optional JAX backend for MMDAgg (`mmdagg(X, Y, backend="jax")`, also accepted by `kt.two_sample`):
+
+```bash
+pip install -e .[jax]
+```
+
+It uses the same random draws as the NumPy path and gives identical decisions; see the *Backends* section of the user guide and `experiments/jax_backend/results.md` for timings.
+
+## Quickstart
+
+Two-sample test (are `X` and `Y` drawn from the same distribution?):
+
+```python
+import numpy as np
+import kt
+
+rng = np.random.default_rng(0)
+X = rng.standard_normal((200, 2))
+Y = rng.standard_normal((200, 2)) + 0.5
+res = kt.two_sample(X, Y)
+print(res)          # TestResult(reject=True, method='mmdagg', alpha=0.05)
+```
+
+Goodness-of-fit test (does `X` come from N(0, I), whose score is `-x`?):
+
+```python
+rng = np.random.default_rng(1)
+X = rng.standard_normal((300, 2)) + 0.3
+res = kt.goodness_of_fit(X, score_X=-X)
+print(res.reject, int(res.details["reject"].sum()), "of", len(res.details["reject"]))
+# True 10 of 10
+```
+
+`res.details` holds diagnostics for each kernel and bandwidth (statistic, quantile, p-value, threshold, reject). See the [documentation](docs/index.md).
+
+## Citing
+
+If you use these tests, please cite the original papers:
+
+```bibtex
+@article{schrab2023mmdagg,
+  author  = {Antonin Schrab and Ilmun Kim and M{\'e}lisande Albert and B{\'e}atrice Laurent and Benjamin Guedj and Arthur Gretton},
+  title   = {{MMD} Aggregated Two-Sample Test},
+  journal = {Journal of Machine Learning Research},
+  year    = {2023}, volume = {24}, number = {194}, pages = {1--81},
+  url     = {http://jmlr.org/papers/v24/21-1289.html}
+}
+@inproceedings{schrab2022ksdagg,
+  author    = {Antonin Schrab and Benjamin Guedj and Arthur Gretton},
+  title     = {{KSD} Aggregated Goodness-of-fit Test},
+  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
+  year      = {2022}
+}
+```
+
+The ports are adapted from the MIT-licensed reference packages [antoninschrab/mmdagg](https://github.com/antoninschrab/mmdagg), [antoninschrab/ksdagg](https://github.com/antoninschrab/ksdagg) and [antoninschrab/mmdfuse](https://github.com/antoninschrab/mmdfuse).
+
+## Testing and verification
+
+```bash
+pip install -e .[test]
+pytest --cov=kt --cov-config=.coveragerc
+```
+
+The unit tests check power against clear alternatives, Type I error control under the null (Monte Carlo, within 3 standard errors), determinism for a fixed seed, input validation, and a regression test for the p-value bug. Parity with the references is checked separately against every MMDAgg kernel option, all four weightings, both resampling schemes, and KSDAgg wild-bootstrap and parametric modes. See [docs/verification.md](docs/verification.md). CI runs on Linux, macOS and Windows with Python 3.9 to 3.13.
+
+## Roadmap
+
+- GPU benchmarks for the JAX backend (implemented, untested on GPU)
+- Fast path for the permutation (unequal sample size) MMDAgg calibration
+
+## Contributing
+
+See [CONTRIBUTING.md](CONTRIBUTING.md). Every new port has to match its reference implementation exactly.
+
+## License
+
+MIT. See [LICENSE](LICENSE).

kerneltests-0.0.1.dist-info/RECORD

@@ -0,0 +1,14 @@
+kerneltests-0.0.1.dist-info/licenses/LICENSE,sha256=JPwDKdmCFInu3md_StwL2_VzYZfE4G8kEMj9787yCPc,1173
+kt/__init__.py,sha256=xeWodsAPMjDWyz7YEWFsgjf3FZf9WBN4qzmlld1l900,274
+kt/_common.py,sha256=r1g-vasoIimmJ_OqQJKZ82SxgcgFkBUgkpI6AN2rhvo,982
+kt/agginc.py,sha256=WAkbtuYxsFFLvRxbJFpOmsLgnTF_DsscrjTwZT7nRX8,12362
+kt/api.py,sha256=zM2KdCJV8n0aPp0PQZP5zNk4urekobMd4jPrz-hxf-Q,3151
+kt/ksdagg.py,sha256=lSK3NEEqkxNIr2GhJQKq1Ut52zE8UPm1sS0Y1kiXBx4,5879
+kt/mmdagg.py,sha256=vTC5fdPltARfYqixeiBGykV7Z3SseIC_TN1e6MajRuI,8209
+kt/mmdfuse.py,sha256=YJRRVTlwODSEEAg2-PG9jSL-3T9-D6oj56dcsQHfa_4,6584
+kt/backends/__init__.py,sha256=sYlzOQqwcNijrbNXHw62KGvqZqTLFolggSxgum3MeTc,594
+kt/backends/jax_mmdagg.py,sha256=4vAF5D4OsECnFyXjMPE7-juswISj3iLZVOXWIkUoITQ,3986
+kerneltests-0.0.1.dist-info/METADATA,sha256=TVhBVStvzmyuHxSxizWn7liUR92mYuCtDHdqYH0uo9g,5222
+kerneltests-0.0.1.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+kerneltests-0.0.1.dist-info/top_level.txt,sha256=wqvF_VC4zbcLJ2p-d_hAHzAatbzOdJLHrbk5wUtYGTw,3
+kerneltests-0.0.1.dist-info/RECORD,,

kerneltests-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

kerneltests-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,10 @@
+MIT License
+
+Copyright (c) 2026 Ankith Burki
+Portions adapted from https://github.com/antoninschrab/mmdagg, Copyright (c) 2023 Antonin Schrab (MIT).
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

kerneltests-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+kt

kt/__init__.py

@@ -0,0 +1,5 @@
+"""kerneltests: unified, tested implementations of modern kernel hypothesis tests."""
+from .api import TestResult, two_sample, goodness_of_fit, independence
+
+__version__ = "0.1.0.dev0"
+__all__ = ["TestResult", "two_sample", "goodness_of_fit", "independence", "__version__"]

kt/_common.py

@@ -0,0 +1,21 @@
+"""Shared helpers. Weights follow Schrab et al. (MIT reference code) in the same
+floating-point evaluation order, which keeps results bit-identical to the references."""
+import numpy as np
+
+WEIGHTS = ("uniform", "decreasing", "increasing", "centred")
+
+
+def weights(N: int, weights_type: str) -> np.ndarray:
+    if weights_type not in WEIGHTS:
+        raise ValueError(f"weights_type must be one of {WEIGHTS}")
+    if weights_type == "uniform":
+        return np.array([1 / N] * N)
+    if weights_type == "decreasing":
+        z = sum(1 / k for k in range(1, N + 1))
+        return np.array([1 / (k * z) for k in range(1, N + 1)])
+    if weights_type == "increasing":
+        z = sum(1 / k for k in range(1, N + 1))
+        return np.array([1 / ((N + 1 - k) * z) for k in range(1, N + 1)])
+    off = 1 if N % 2 == 1 else 0.5
+    z = sum(1 / (abs((N + 1) / 2 - k) + off) for k in range(1, N + 1))
+    return np.array([1 / ((abs((N + 1) / 2 - k) + off) * z) for k in range(1, N + 1)])

kt/agginc.py

@@ -0,0 +1,264 @@
+"""MMDAggInc, HSICAggInc and KSDAggInc: efficient aggregated kernel tests using
+incomplete U-statistics (Schrab, Kim, Guedj, Gretton; NeurIPS 2022).
+
+Re-implementation adapted from the MIT-licensed reference package
+https://github.com/antoninschrab/agginc (Copyright (c) 2023 Antonin Schrab),
+NumPy version ``agginc/np.py``. Same random-number protocol as the reference
+(one ``RandomState(seed).choice([1., -1.], size=(N, B))`` Rademacher draw), so
+results agree exactly for a given seed. Index construction, kernel evaluation
+across bandwidths and the quantile bisection are vectorized; inputs are
+validated with ValueError. Fixed relative to the reference: missing
+``import warnings`` (NameError on non-default ``number_bandwidths`` for hsic
+or non-default ``hsic_collection_parameters`` for mmd/ksd), and the custom
+HSIC bandwidth check testing ``len(bandwidths[0])`` twice instead of
+``bandwidths[1]``.
+"""
+from __future__ import annotations
+import warnings
+import numpy as np
+from scipy.spatial.distance import cdist, pdist
+
+_TYPES = ("mmd", "hsic", "ksd")
+_WEIGHTS = ("uniform", "decreasing", "increasing", "centred")
+
+
+def _weights(N, kind):
+    i = np.arange(1, N + 1)
+    if kind == "uniform":
+        return np.array([1 / N] * N)
+    if kind in ("decreasing", "increasing"):
+        h = sum(1 / k for k in range(1, N + 1))
+        return 1 / (i * h) if kind == "decreasing" else 1 / ((N + 1 - i) * h)
+    if kind == "centred":
+        c = np.abs((N + 1) / 2 - i) + (1 if N % 2 == 1 else 0.5)
+        return 1 / (c * sum(1 / c))
+    raise ValueError(f"weights_type must be one of {_WEIGHTS}")
+
+
+def _indices(N, R):
+    r = np.arange(1, R + 1)
+    lens = np.maximum(N - r, 0)
+    ii = np.concatenate([np.arange(L) for L in lens]) if lens.sum() else np.zeros(0, int)
+    jj = ii + np.repeat(r, lens)
+    return ii, jj
+
+
+def _h_mmd(X, Y, R, bws):
+    N = min(X.shape[0], Y.shape[0])
+    ii, jj = _indices(N, R)
+    sq = [np.linalg.norm(A[ii] - B[jj], axis=1) ** 2 for A, B in ((X, X), (X, Y), (Y, X), (Y, Y))]
+    b2 = (np.asarray(bws, dtype=float) ** 2)[:, None]
+    k = [np.exp(-s[None, :] / b2) for s in sq]
+    return k[0] - k[1] - k[2] + k[3], ii, jj, N
+
+
+def _h_hsic(X, Y, R, bws):
+    N = X.shape[0] // 2
+    hx, ii, jj, _ = _h_mmd(X[:N], X[N:], R, bws[0])
+    hy = _h_mmd(Y[:N], Y[N:], R, bws[1])[0]
+    return (hx[None] * hy[:, None]).reshape(hx.shape[0] * hy.shape[0], -1), ii, jj, N
+
+
+def _h_ksd(X, S, R, bws, beta=-0.5):
+    N, d = X.shape
+    ii, jj = _indices(N, R)
+    D = X[ii] - X[jj]
+    nrm = np.linalg.norm(D, axis=1) ** 2
+    si, sj = S[ii], S[jj]
+    a = np.einsum("ij,ij->i", si - sj, D, optimize=True)
+    c = np.einsum("ij,ij->i", si, sj, optimize=True)
+    h = np.zeros((len(bws), len(ii)))
+    for r, b in enumerate(bws):  # kept as a loop: identical op order to the reference
+        bn = b ** 2 + nrm
+        h[r] = (c * bn ** beta - a * 2 * beta * bn ** (beta - 1)
+                - 2 * d * beta * bn ** (beta - 1) - 4 * beta * (beta - 1) * nrm * bn ** (beta - 2))
+    return h, ii, jj, N
+
+
+_H = {"mmd": _h_mmd, "hsic": _h_hsic, "ksd": _h_ksd}
+
+
+def _bootstrap(h, ii, jj, N, B, seed, batch_size, memory_percentage=0.8):
+    eps = np.random.RandomState(seed).choice([1.0, -1.0], size=(N, B))
+    if batch_size is None:
+        bs = B
+    elif batch_size == "auto":
+        # Reference bug: it divides RAM by the bytes of a length-#bandwidths
+        # vector, so "auto" is always one batch of size len(ii) x B (OOM for
+        # large n). Size by the real per-column cost instead; batching does not
+        # change any result.
+        # psutil can over-report memory inside containers, so cap at ~256 MB per batch.
+        try:
+            import psutil
+            mem = min(psutil.virtual_memory().available, 2 ** 28 / memory_percentage)
+        except Exception:  # pragma: no cover
+            mem = 2 ** 28
+        bs = int(max(1, min(B, memory_percentage * mem / (3 * 8 * max(len(ii), 1)))))
+    else:
+        bs = batch_size
+    out = np.zeros((h.shape[0], B))
+    for s in range(0, B, bs):
+        e = eps[:, s:s + bs]
+        out[:, s:s + bs] = h @ (e[ii] * e[jj])
+    return out / len(ii), h @ np.ones(h.shape[1]) / len(ii)
+
+
+def _median(D):
+    return np.median(D[D > 0])
+
+
+def _check_common(agg_type, X, Y, R, alpha, batch_size, memory_percentage, seed):
+    if agg_type not in _TYPES:
+        raise ValueError('agg_type should be "mmd" or "hsic" or "ksd".')
+    X = np.asarray(X, dtype=float); Y = np.asarray(Y, dtype=float)
+    if X.ndim != 2 or Y.ndim != 2:
+        raise ValueError("X and Y must be 2-D arrays")
+    if agg_type == "mmd" and X.shape[1] != Y.shape[1]:
+        raise ValueError("mmd: X and Y must have the same number of columns")
+    if agg_type == "hsic" and X.shape[0] != Y.shape[0]:
+        raise ValueError("hsic: X and Y must have the same number of rows")
+    if agg_type == "ksd" and X.shape != Y.shape:
+        raise ValueError("ksd: Y (score of X) must have the same shape as X")
+    n = min(X.shape[0], Y.shape[0]) // (2 if agg_type == "hsic" else 1)
+    if n < 2:
+        raise ValueError("too few samples (need >= 2 effective samples; hsic needs >= 4 rows)")
+    if not (isinstance(R, (int, np.integer)) and R >= 1):
+        raise ValueError("R must be a positive int")
+    if not 0 < alpha < 1:
+        raise ValueError("alpha must be in (0, 1)")
+    if not (batch_size is None or batch_size == "auto"
+            or (isinstance(batch_size, (int, np.integer)) and batch_size > 0)):
+        raise ValueError('batch_size must be None, "auto" or a positive int')
+    if not 0 < memory_percentage <= 1:
+        raise ValueError("memory_percentage must be in (0, 1]")
+    return X, Y
+
+
+def agginc(agg_type, X, Y, R=200, alpha=0.05, weights_type="uniform", number_bandwidths=10,
+           hsic_collection_parameters=(2, -2, 2), batch_size="auto", memory_percentage=0.8,
+           B1=500, B2=500, B3=50, seed=42, return_details=False, bandwidths=None):
+    """Aggregated incomplete-U test. Returns 1 if H0 is rejected, else 0.
+
+    agg_type "mmd": two-sample (X, Y); "hsic": independence (paired X, Y);
+    "ksd": goodness-of-fit (Y is the score of the model at X). With
+    return_details=True returns (output, dict of per-test arrays).
+    """
+    X, Y = _check_common(agg_type, X, Y, R, alpha, batch_size, memory_percentage, seed)
+    for name, v in (("B1", B1), ("B2", B2), ("B3", B3)):
+        if not (isinstance(v, (int, np.integer)) and v > 0):
+            raise ValueError(f"{name} must be a positive int")
+    if weights_type not in _WEIGHTS:
+        raise ValueError(f"weights_type must be one of {_WEIGHTS}")
+    if number_bandwidths != 10 and agg_type == "hsic":
+        warnings.warn("number_bandwidths has no effect for hsic; use hsic_collection_parameters or bandwidths.")
+    if agg_type != "hsic" and tuple(hsic_collection_parameters) != (2, -2, 2):
+        warnings.warn("hsic_collection_parameters has no effect for mmd/ksd; use number_bandwidths or bandwidths.")
+
+    if bandwidths is not None:
+        if agg_type == "hsic":
+            if len(bandwidths) != 2 or len(bandwidths[0]) == 0 or len(bandwidths[1]) == 0:
+                raise ValueError("hsic bandwidths must be a pair of non-empty sequences")
+            bandwidths = [list(bandwidths[0]), list(bandwidths[1])]
+        else:
+            bandwidths = np.asarray(bandwidths, dtype=float)
+            if bandwidths.ndim != 1 or bandwidths.size == 0:
+                raise ValueError("bandwidths must be a non-empty 1-D array")
+    elif agg_type == "hsic":
+        if weights_type != "uniform":
+            raise ValueError("hsic requires weights_type='uniform'")
+        p = tuple(hsic_collection_parameters)
+        if len(p) == 3:
+            power, lm, lp = p; lm, lp = (lm, lm), (lp, lp)
+        elif len(p) == 5:
+            power, lmx, lpx, lmy, lpy = p; lm, lp = (lmx, lmy), (lpx, lpy)
+        else:
+            raise ValueError("hsic_collection_parameters should have length either 3 or 5.")
+        med = (_median(pdist(X[:500], "euclidean")), _median(pdist(Y[:500], "euclidean")))
+        bandwidths = [[power ** i * med[j] for i in range(lm[j], lp[j] + 1)] for j in range(2)]
+        if not bandwidths[0] or not bandwidths[1]:
+            raise ValueError("hsic_collection_parameters give an empty bandwidth collection")
+    else:
+        if not (isinstance(number_bandwidths, (int, np.integer)) and number_bandwidths > 1):
+            raise ValueError("number_bandwidths must be an int > 1")
+        if agg_type == "mmd":
+            D = cdist(X[:500], Y[:500], "euclidean").reshape(-1)
+            if np.min(D) < 1e-1:
+                dd = np.sort(D); lmin = np.maximum(dd[int(np.floor(len(dd) * 0.05))], 1e-1)
+            else:
+                lmin = np.min(D)
+            lmin = lmin / 2
+            lmax = np.maximum(np.max(D), 3e-1) * 2
+            pw = (lmax / lmin) ** (1 / (number_bandwidths - 1))
+            bandwidths = np.array([pw ** i * lmin for i in range(number_bandwidths)])
+        else:
+            D = pdist(X[:500], "euclidean"); D = D[D > 0]
+            lmax = np.maximum(np.max(D), 2)
+            pw = (lmax / 1) ** (1 / (number_bandwidths - 1))
+            bandwidths = np.array([pw ** i * 1 / X.shape[1] for i in range(number_bandwidths)])
+
+    h, ii, jj, N = _H[agg_type](X, Y, R, bandwidths)
+    boot, stat = _bootstrap(h, ii, jj, N, B1 + B2, seed, batch_size, memory_percentage)
+    nb = boot.shape[0]
+    w = _weights(nb, weights_type)
+    b1 = np.column_stack([boot[:, :B1], stat])
+    b1s = np.sort(b1); b2 = boot[:, B1:]
+    rows = np.arange(nb)
+
+    def quantiles(u):
+        return b1s[rows, np.ceil((B1 + 1) * (1 - u * w)).astype(int) - 1]
+
+    lo, hi = 0.0, np.min(1 / w)
+    for _ in range(B3):
+        u = (hi + lo) / 2
+        if np.sum(np.max(b2 - quantiles(u)[:, None], 0) > 0) / B2 <= alpha:
+            lo = u
+        else:
+            hi = u
+    q = quantiles(lo)
+    pv = np.mean(b1 - stat[:, None] >= 0, -1)
+    rej = stat > q
+    out = int(rej.any())
+    if not return_details:
+        return out
+    if agg_type == "hsic":
+        bwx = np.array([bandwidths[0][i] for j in range(len(bandwidths[1])) for i in range(len(bandwidths[0]))])
+        bwy = np.array([bandwidths[1][j] for j in range(len(bandwidths[1])) for i in range(len(bandwidths[0]))])
+        bwinfo = {"bandwidth_x": bwx, "bandwidth_y": bwy}
+    else:
+        bwinfo = {"bandwidth": np.asarray(bandwidths)}
+    return out, {"kernel": "imq" if agg_type == "ksd" else "gaussian", **bwinfo, "statistic": stat,
+                 "quantile": q, "p_value": pv, "threshold": lo * w, "reject": rej}
+
+
+def inc(agg_type, X, Y, R=200, alpha=0.05, batch_size="auto", memory_percentage=0.8, B=500,
+        seed=42, return_details=False, bandwidth=None):
+    """Single-bandwidth incomplete-U test (median heuristic unless bandwidth given)."""
+    X, Y = _check_common(agg_type, X, Y, R, alpha, batch_size, memory_percentage, seed)
+    if not (isinstance(B, (int, np.integer)) and B > 0):
+        raise ValueError("B must be a positive int")
+    if bandwidth is not None:
+        if agg_type == "hsic":
+            if len(bandwidth) != 2:
+                raise ValueError("hsic bandwidth must be a pair (bandwidth_X, bandwidth_Y)")
+            bws = (np.array(bandwidth[0]).reshape(1), np.array(bandwidth[1]).reshape(1))
+        else:
+            bws = np.array(bandwidth).reshape(1)
+    elif agg_type == "mmd":
+        bws = np.array([_median(cdist(X[:500], Y[:500], "euclidean").reshape(-1))])
+    elif agg_type == "hsic":
+        bws = (np.array([_median(pdist(X[:500], "euclidean"))]), np.array([_median(pdist(Y[:500], "euclidean"))]))
+    else:
+        bws = np.array([_median(pdist(X[:500], "euclidean"))])
+    h, ii, jj, N = _H[agg_type](X, Y, R, bws)
+    boot, stat = _bootstrap(h, ii, jj, N, B, seed, batch_size, memory_percentage)
+    stat = stat[0]
+    b1 = np.column_stack([boot, stat])
+    q = np.sort(b1)[0, int(np.ceil((B + 1) * (1 - alpha))) - 1]
+    pv = np.mean(b1 - stat >= 0, -1)[0]
+    out = int(stat > q)
+    if not return_details:
+        return out
+    bwinfo = ({"bandwidth_x": bws[0][0], "bandwidth_y": bws[1][0]} if agg_type == "hsic"
+              else {"bandwidth": bws[0]})
+    return out, {"kernel": "imq" if agg_type == "ksd" else "gaussian", **bwinfo, "statistic": stat,
+                 "quantile": q, "p_value": pv, "threshold": alpha, "reject": bool(stat > q)}

kt/api.py

@@ -0,0 +1,74 @@
+"""Unified entry points: one function per testing problem, one result type."""
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class TestResult:
+    """Outcome of a kernel hypothesis test.
+
+    reject  -- True if the null hypothesis is rejected at level ``alpha``.
+    method  -- name of the test that was run.
+    alpha   -- significance level used.
+    details -- per-kernel/bandwidth diagnostics (statistics, quantiles, p-values, ...).
+    """
+    reject: bool
+    method: str
+    alpha: float
+    details: Dict[str, Any] = field(default_factory=dict, repr=False)
+
+    def __bool__(self) -> bool:
+        return self.reject
+
+
+def _run(fn: Callable, method: str, alpha: float, *args, **kwargs) -> TestResult:
+    out, details = fn(*args, alpha=alpha, return_details=True, **kwargs)
+    return TestResult(reject=bool(out), method=method, alpha=alpha, details=details)
+
+
+def two_sample(X, Y, method: str = "mmdagg", alpha: float = 0.05, **kwargs) -> TestResult:
+    """Test H0: X and Y are drawn from the same distribution.
+
+    method: "mmdagg" (Schrab et al., JMLR 2023; quadratic time) or
+            "mmdagginc" (Schrab et al., NeurIPS 2022; incomplete U-statistics, faster for large n) or
+            "mmdfuse" (Biggs, Schrab, Gretton, NeurIPS 2023; fused soft-max over kernels, permutations).
+    """
+    if method == "mmdagg":
+        from .mmdagg import mmdagg
+        return _run(mmdagg, method, alpha, X, Y, **kwargs)
+    if method == "mmdagginc":
+        from .agginc import agginc
+        return _run(lambda *a, **k: agginc("mmd", *a, **k), method, alpha, X, Y, **kwargs)
+    if method == "mmdfuse":
+        from .mmdfuse import mmdfuse
+        return _run(mmdfuse, method, alpha, X, Y, **kwargs)
+    raise ValueError(f"unknown two-sample method {method!r}; available: 'mmdagg', 'mmdagginc', 'mmdfuse'")
+
+
+def independence(X, Y, method: str = "hsicagginc", alpha: float = 0.05, **kwargs) -> TestResult:
+    """Test H0: paired samples X[i], Y[i] come from independent variables.
+
+    method: "hsicagginc" (Schrab et al., NeurIPS 2022).
+    """
+    if method == "hsicagginc":
+        from .agginc import agginc
+        return _run(lambda *a, **k: agginc("hsic", *a, **k), method, alpha, X, Y, **kwargs)
+    raise ValueError(f"unknown independence method {method!r}; available: 'hsicagginc'")
+
+
+def goodness_of_fit(X, score_X, method: str = "ksdagg", alpha: float = 0.05, **kwargs) -> TestResult:
+    """Test H0: X is drawn from the model whose score function values at X are ``score_X``.
+
+    method: "ksdagg" (Schrab, Guedj, Gretton, NeurIPS 2022) or "ksdagginc" (incomplete U-statistics).
+    """
+    if method == "ksdagg":
+        from .ksdagg import ksdagg
+        X = np.asarray(X, dtype=float)
+        return _run(ksdagg, method, alpha, X, score_X, **kwargs)
+    if method == "ksdagginc":
+        from .agginc import agginc
+        return _run(lambda *a, **k: agginc("ksd", *a, **k), method, alpha, X, score_X, **kwargs)
+    raise ValueError(f"unknown goodness-of-fit method {method!r}; available: 'ksdagg', 'ksdagginc'")

kt/backends/__init__.py

@@ -0,0 +1,15 @@
+"""Optional compute backends. Each backend module is imported lazily."""
+
+
+def require_jax():
+    """Import and return ``jax`` (with float64 enabled) or raise a helpful ImportError."""
+    try:
+        import jax
+    except ImportError as e:
+        raise ImportError(
+            "backend='jax' requires the optional dependency JAX. Install it with "
+            "`pip install kerneltests[jax]` (or `pip install \"jax[cpu]\"`; see "
+            "https://jax.readthedocs.io for GPU wheels), or use backend='numpy'."
+        ) from e
+    jax.config.update("jax_enable_x64", True)
+    return jax

kt/backends/jax_mmdagg.py

@@ -0,0 +1,107 @@
+"""JAX backend for MMDAgg: jit-compiled kernel matrices and bootstrap/permutation statistics.
+
+The random draws (Rademacher matrix R or permutation matrices V) are generated by the
+NumPy ``RandomState`` protocol in :func:`kt.mmdagg.mmdagg` and passed in, so both
+backends see identical resamples. Distances, kernel matrices (all bandwidths of one
+kernel batched with ``jax.vmap``) and the quadratic forms r^T K r are computed in JAX
+in float64. Runs on whatever device JAX selects (CPU, or GPU with a CUDA jaxlib).
+"""
+from __future__ import annotations
+from functools import partial
+
+import numpy as np
+
+from . import require_jax
+
+jax = require_jax()
+import jax.numpy as jnp  # noqa: E402
+
+_MEM_BUDGET = 1 << 30  # bytes of kernel matrices materialised per vmap batch
+
+
+def _kernel(d, base):
+    if base == "gaussian":
+        return jnp.exp(-d ** 2)
+    if base == "imq":
+        return (1 + d ** 2) ** (-0.5)
+    if base == "matern_0.5":
+        return jnp.exp(-d)
+    if base == "matern_1.5":
+        return (1 + np.sqrt(3) * d) * jnp.exp(-np.sqrt(3) * d)
+    if base == "matern_2.5":
+        return (1 + np.sqrt(5) * d + 5 / 3 * d ** 2) * jnp.exp(-np.sqrt(5) * d)
+    if base == "matern_3.5":
+        return (1 + np.sqrt(7) * d + 2 * 7 / 5 * d ** 2 + 7 * np.sqrt(7) / 3 / 5 * d ** 3) * jnp.exp(-np.sqrt(7) * d)
+    if base == "matern_4.5":
+        return (1 + 3 * d + 3 * (6 ** 2) / 28 * d ** 2 + (6 ** 3) / 84 * d ** 3
+                + (6 ** 4) / 1680 * d ** 4) * jnp.exp(-3 * d)
+    raise ValueError(f"unknown kernel {base!r}")
+
+
+def _base(kernel):
+    if kernel in ("gaussian", "imq"):
+        return kernel
+    return "matern_0.5" if kernel == "laplace" else kernel[:10]
+
+
+@partial(jax.jit, static_argnames="metric")
+def _pairwise(A, B, metric):
+    diff = A[:, None, :] - B[None, :, :]
+    if metric == "cityblock":
+        return jnp.sum(jnp.abs(diff), -1)
+    return jnp.sqrt(jnp.sum(diff ** 2, -1))
+
+
+@partial(jax.jit, static_argnames="base")
+def _wild(Pxx, Pyy, Pxy, bws, Rn, base):
+    n = Pxx.shape[0]
+    off = 1.0 - jnp.eye(n)
+
+    def one(bw):
+        H = (_kernel(Pxx / bw, base) + _kernel(Pyy / bw, base)
+             - _kernel(Pxy / bw, base) - _kernel(Pxy.T / bw, base)) * off
+        return jnp.sum(Rn * (H @ Rn), 0)
+    return jax.vmap(one)(bws)
+
+
+@partial(jax.jit, static_argnames=("base", "m", "n"))
+def _perm(P, bws, V10, V01, V11, base, m, n):
+    off = 1.0 - jnp.eye(P.shape[0])
+
+    def one(bw):
+        K = _kernel(P / bw, base) * off
+        return (jnp.sum(V10 * (K @ V10), 0) * (n - m + 1) / (m * n * (m - 1))
+                + jnp.sum(V01 * (K @ V01), 0) * (m - n + 1) / (m * n * (n - 1))
+                + jnp.sum(V11 * (K @ V11), 0) / (m * n))
+    return jax.vmap(one)(bws)
+
+
+def statistics(X, Y, kbl, wild, draws):
+    """Return the (#kernels*#bandwidths, B1+B2+1) matrix of bootstrap/permutation statistics.
+
+    kbl: list of (kernel name, bandwidths, metric); draws: (Rn,) if wild else (V10, V01, V11).
+    """
+    m, n = X.shape[0], Y.shape[0]
+    Xj, Yj = jnp.asarray(X), jnp.asarray(Y)
+    side = n if wild else m + n
+    per_bw = (4 if wild else 1) * side * side * 8
+    batch = max(1, _MEM_BUDGET // per_bw)
+    dev = [jnp.asarray(a) for a in draws]
+    cache, rows = {}, []
+    for kern, bandwidths, metric in kbl:
+        if metric not in cache:
+            if wild:
+                cache = {metric: (_pairwise(Xj, Xj, metric), _pairwise(Yj, Yj, metric),
+                                  _pairwise(Xj, Yj, metric))}
+            else:
+                Z = jnp.vstack((Xj, Yj))
+                cache = {metric: (_pairwise(Z, Z, metric),)}
+        P = cache[metric]
+        bws = jnp.asarray(bandwidths, dtype=jnp.float64)
+        for s in range(0, len(bandwidths), batch):
+            b = bws[s:s + batch]
+            if wild:
+                rows.append(_wild(*P, b, dev[0], base=_base(kern)))
+            else:
+                rows.append(_perm(P[0], b, *dev, base=_base(kern), m=m, n=n))
+    return np.asarray(jnp.concatenate(rows, 0))

kt/ksdagg.py

@@ -0,0 +1,129 @@
+"""KSDAgg goodness-of-fit test (Schrab, Guedj, Gretton; NeurIPS 2022).
+
+Re-implementation adapted from the MIT-licensed reference package
+https://github.com/antoninschrab/ksdagg (Copyright (c) 2023 Antonin Schrab).
+Same random-number protocol and floating-point formulas as the reference
+(``ksdagg/np.py``), so results agree exactly for a given seed; the Stein
+kernel shares its bandwidth-independent terms, and the bisection step and
+quantile lookup are vectorized.
+
+Supported: kernel "imq" (the reference's only kernel), approx_type
+"wild bootstrap" and "parametric", all four weights_type values, and a
+user-supplied bandwidth collection. Differences from the reference:
+ValueError validation instead of asserts; ``bandwidths`` may be any 1-D
+array-like (the reference silently ignores non-ndarray inputs); for
+"parametric", B1/B2 must have shape (number_bandwidths, B) -- the shape the
+reference code actually uses, not the transposed shape in its docstring --
+and number_bandwidths is taken from B1.shape[0] when bandwidths is None.
+"""
+from __future__ import annotations
+import numpy as np
+from scipy.spatial.distance import pdist
+from ._common import WEIGHTS as _WEIGHTS, weights as _weights
+
+
+
+def _stein_imq(X, S, bandwidths):
+    """Yield the zero-diagonal IMQ Stein kernel matrix (beta = 1/2) per bandwidth."""
+    beta = 0.5
+    p = X.shape[1]
+    norms = np.sum(X ** 2, -1)
+    dists = -2 * X @ X.T + norms[:, None] + norms[None, :]
+    diffs = np.sum(X * S, -1)[:, None] - X @ S.T
+    diffs = diffs + diffs.T
+    SS = S @ S.T
+    for bw in bandwidths:
+        g = 1 / bw ** 2
+        res = 1 + g * dists
+        r1 = res ** (-beta - 1)
+        H = (SS * res ** (-beta) + 2 * beta * g * r1 * diffs
+             + 2 * (beta * g * r1 * p
+                    - 2 * beta * (beta + 1) * g ** 2 * dists * res ** (-beta - 2)))
+        np.fill_diagonal(H, 0)
+        yield H
+
+
+def ksdagg(X, score_X, alpha=0.05, number_bandwidths=10, weights_type="uniform",
+           approx_type="wild bootstrap", kernel="imq", B1=2000, B2=2000, B3=50,
+           seed=42, bandwidths=None, return_details=False):
+    """Return 1 if KSDAgg rejects H0: X ~ model (whose score at X is score_X), else 0."""
+    X = np.asarray(X, dtype=float); S = np.asarray(score_X, dtype=float)
+    if X.ndim != 2 or X.shape != S.shape:
+        raise ValueError("X and score_X must be 2-D arrays of the same shape")
+    m = X.shape[0]
+    if m < 2:
+        raise ValueError("need at least 2 samples")
+    if not 0 < alpha < 1:
+        raise ValueError("alpha must be in (0, 1)")
+    if kernel != "imq":
+        raise ValueError(f"unsupported kernel {kernel!r}; only 'imq'")
+    if weights_type not in _WEIGHTS:
+        raise ValueError(f"weights_type must be one of {_WEIGHTS}")
+    if not (isinstance(B3, int) and B3 > 0):
+        raise ValueError("B3 must be a positive int")
+    if approx_type == "wild bootstrap":
+        for name, v in (("B1", B1), ("B2", B2)):
+            if not (isinstance(v, int) and v > 0):
+                raise ValueError(f"{name} must be a positive int")
+        P1 = P2 = None
+    elif approx_type == "parametric":
+        P1 = np.asarray(B1, dtype=float); P2 = np.asarray(B2, dtype=float)
+        if P1.ndim != 2 or P2.ndim != 2 or P1.shape[0] != P2.shape[0] or min(P1.shape + P2.shape) < 1:
+            raise ValueError("parametric B1, B2 must be non-empty 2-D arrays (number_bandwidths, B)")
+        B1, B2 = P1.shape[1], P2.shape[1]
+        if bandwidths is None:
+            number_bandwidths = P1.shape[0]
+    else:
+        raise ValueError("approx_type must be 'wild bootstrap' or 'parametric'")
+
+    if bandwidths is not None:
+        bandwidths = np.asarray(bandwidths, dtype=float)
+        if bandwidths.ndim != 1 or len(bandwidths) < 1 or np.any(bandwidths <= 0):
+            raise ValueError("bandwidths must be a non-empty 1-D array of positive values")
+        number_bandwidths = len(bandwidths)
+    if not (isinstance(number_bandwidths, int) and number_bandwidths > 1) and bandwidths is None:
+        raise ValueError("number_bandwidths must be an int > 1")
+    if P1 is not None and P1.shape[0] != number_bandwidths:
+        raise ValueError("parametric B1/B2 must have one row per bandwidth")
+    if bandwidths is None:
+        d = pdist(X[:500], "euclidean"); d = d[d > 0]
+        lam_max = np.maximum(np.max(d), 2) if d.size else 2.0
+        power = (lam_max / 1) ** (1 / (number_bandwidths - 1))
+        bandwidths = np.array([power ** i * 1 / X.shape[1] for i in range(number_bandwidths)])
+    N = number_bandwidths
+    w = _weights(N, weights_type)
+
+    if P1 is None:
+        rs = np.random.RandomState(seed)
+        R = rs.choice([1.0, -1.0], size=(B1 + B2 + 1, m)); R[B1] = 1.0
+        R = R.T
+        M = np.empty((N, B1 + B2 + 1))
+        for i, H in enumerate(_stein_imq(X, S, bandwidths)):
+            M[i] = np.sum(R * (H @ R), 0) / (m * (m - 1))
+        stat = M[:, B1]
+        M1 = np.sort(M[:, :B1 + 1]); M2 = M[:, B1 + 1:]
+    else:
+        r = np.ones(m)
+        stat = np.array([r @ H @ r / (m * (m - 1)) for H in _stein_imq(X, S, bandwidths)])
+        M1 = np.sort(np.concatenate((P1, stat[:, None]), axis=1)); M2 = P2
+
+    rows = np.arange(N)
+    def quantiles(u):
+        return M1[rows, np.ceil((B1 + 1) * (1 - u * w)).astype(int) - 1]
+
+    u_lo, u_hi = 0, np.min(1 / w)
+    for _ in range(B3):
+        u = (u_hi + u_lo) / 2
+        if np.sum(np.max(M2 - quantiles(u)[:, None], 0) > 0) / B2 <= alpha:
+            u_lo = u
+        else:
+            u_hi = u
+    u = u_lo
+    q = quantiles(u)
+    reject_each = stat > q
+    out = int(reject_each.any())
+    if not return_details:
+        return out
+    pvals = np.mean(M1 - stat[:, None] >= 0, -1)
+    return out, {"kernel": ["imq"] * N, "bandwidth": np.asarray(bandwidths, float), "ksd": stat,
+                 "quantile": q, "p_value": pvals, "threshold": u * w, "reject": reject_each}

kt/mmdagg.py

@@ -0,0 +1,177 @@
+"""MMDAgg two-sample test (Schrab, Kim, Albert, Laurent, Guedj, Gretton; JMLR 2023).
+
+Re-implementation adapted from the MIT-licensed reference package
+https://github.com/antoninschrab/mmdagg (Copyright (c) 2023 Antonin Schrab).
+Same random-number protocol as the reference, so results agree exactly for a
+given seed; the bisection step and quantile lookup are vectorized, and the
+per-test p-value reporting bug of the reference (p_vals[i] vs p_vals[idx]) is fixed.
+"""
+from __future__ import annotations
+import numpy as np
+from scipy.spatial.distance import cdist
+from ._common import WEIGHTS as _WEIGHTS, weights as _weights
+
+_L1 = ("laplace",) + tuple(f"matern_{i}.5_l1" for i in range(5))
+_L2 = ("gaussian", "imq") + tuple(f"matern_{i}.5_l2" for i in range(5))
+
+
+def _bandwidths(distances: np.ndarray, number: int) -> np.ndarray:
+    if np.min(distances) < 1e-1:
+        d = np.sort(distances)
+        lam_min = np.maximum(d[int(np.floor(len(d) * 0.05))], 1e-1)
+    else:
+        lam_min = np.min(distances)
+    lam_min = lam_min / 2
+    lam_max = np.maximum(np.max(distances), 3e-1) * 2
+    power = (lam_max / lam_min) ** (1 / (number - 1))
+    return np.array([power ** i * lam_min for i in range(number)])
+
+
+def _kernel(pairwise: np.ndarray, kernel: str, bandwidth: float) -> np.ndarray:
+    d = pairwise / bandwidth
+    if kernel == "gaussian":
+        return np.exp(-d ** 2)
+    if kernel == "imq":
+        return (1 + d ** 2) ** (-0.5)
+    base = "matern_0.5" if kernel == "laplace" else kernel[:10]
+    if base == "matern_0.5":
+        return np.exp(-d)
+    if base == "matern_1.5":
+        return (1 + np.sqrt(3) * d) * np.exp(-np.sqrt(3) * d)
+    if base == "matern_2.5":
+        return (1 + np.sqrt(5) * d + 5 / 3 * d ** 2) * np.exp(-np.sqrt(5) * d)
+    if base == "matern_3.5":
+        return (1 + np.sqrt(7) * d + 2 * 7 / 5 * d ** 2 + 7 * np.sqrt(7) / 3 / 5 * d ** 3) * np.exp(-np.sqrt(7) * d)
+    if base == "matern_4.5":
+        return (1 + 3 * d + 3 * (6 ** 2) / 28 * d ** 2 + (6 ** 3) / 84 * d ** 3 + (6 ** 4) / 1680 * d ** 4) * np.exp(-3 * d)
+    raise ValueError(f"unknown kernel {kernel!r}")
+
+
+def mmdagg(X, Y, alpha=0.05, kernel="laplace_gaussian", number_bandwidths=10,
+           weights_type="uniform", B1=2000, B2=2000, B3=50, seed=42, permutations_same_sample_size=False,
+           return_details=False, fast=True, backend="numpy"):
+    """Return 1 if MMDAgg rejects H0: P_X = P_Y at level alpha, else 0.
+
+    fast: for equal sample sizes (wild bootstrap), evaluate r^T H r with the n x n matrix
+        H = Kxx + Kyy - Kxy - Kyx instead of the reference's 2n x 2n product (about 4x fewer
+        flops; mathematically identical, equal to the reference up to floating-point rounding).
+        Set fast=False to reproduce the reference bit-for-bit.
+    backend: "numpy" (default) or "jax". The JAX backend (optional dependency,
+        ``pip install kerneltests[jax]``) builds all kernel matrices and resampled statistics in
+        jit-compiled float64 code on the default JAX device (CPU or GPU). Random draws are the
+        same NumPy ones, so results match backend="numpy" up to floating-point rounding.
+        ``fast`` is ignored by the JAX backend (it always uses the n x n wild-bootstrap form).
+    """
+    if backend not in ("numpy", "jax"):
+        raise ValueError("backend must be 'numpy' or 'jax'")
+    if backend == "jax":
+        from .backends.jax_mmdagg import statistics as _jax_statistics
+    X = np.asarray(X, dtype=float); Y = np.asarray(Y, dtype=float)
+    if X.ndim != 2 or Y.ndim != 2 or X.shape[1] != Y.shape[1]:
+        raise ValueError("X and Y must be 2-D arrays with the same number of columns")
+    m, n = X.shape[0], Y.shape[0]
+    if m < 2 or n < 2:
+        raise ValueError("need at least 2 samples in each of X and Y")
+    if not 0 < alpha < 1:
+        raise ValueError("alpha must be in (0, 1)")
+    if not (isinstance(number_bandwidths, int) and number_bandwidths > 1):
+        raise ValueError("number_bandwidths must be an int > 1")
+    for name, v in (("B1", B1), ("B2", B2), ("B3", B3)):
+        if not (isinstance(v, int) and v > 0):
+            raise ValueError(f"{name} must be a positive int")
+    wild = (m == n) and not permutations_same_sample_size
+
+    ms = 500
+    bw_l1 = _bandwidths(cdist(X[:ms], Y[:ms], "cityblock").ravel(), number_bandwidths)
+    bw_l2 = _bandwidths(cdist(X[:ms], Y[:ms], "euclidean").ravel(), number_bandwidths)
+    if kernel in _L1:
+        kbl = [(kernel, bw_l1, "cityblock")]
+    elif kernel in _L2:
+        kbl = [(kernel, bw_l2, "euclidean")]
+    elif kernel in ("laplace_gaussian", "gaussian_laplace"):
+        kbl = [("laplace", bw_l1, "cityblock"), ("gaussian", bw_l2, "euclidean")]
+    elif kernel == "all_matern_l1":
+        kbl = [(f"matern_{i}.5_l1", bw_l1, "cityblock") for i in range(5)]
+    elif kernel == "all_matern_l2":
+        kbl = [(f"matern_{i}.5_l2", bw_l2, "euclidean") for i in range(5)]
+    elif kernel in ("all_matern_l1_l2", "all"):
+        kbl = ([(f"matern_{i}.5_l1", bw_l1, "cityblock") for i in range(5)]
+               + [(f"matern_{i}.5_l2", bw_l2, "euclidean") for i in range(5)])
+        if kernel == "all":
+            kbl += [("gaussian", bw_l2, "euclidean"), ("imq", bw_l2, "euclidean")]
+    else:
+        raise ValueError(f"unsupported kernel {kernel!r}")
+    if weights_type not in _WEIGHTS:
+        raise ValueError(f"weights_type must be one of {_WEIGHTS}")
+    w = _weights(number_bandwidths, weights_type) / len(kbl)
+
+    rs = np.random.RandomState(seed)
+    Btot = B1 + B2 + 1
+    if wild:
+        R = rs.choice([-1.0, 1.0], size=(Btot, n)); R[B1] = 1.0
+        Rn = R.T.copy()
+        R = np.concatenate((R.T, -R.T))
+    else:
+        idx = rs.rand(Btot, m + n).argsort(axis=1)
+        def perm(v):
+            V = np.take_along_axis(np.tile(v, (Btot, 1)), idx, axis=1); V[B1] = v; return V.T
+        V11 = perm(np.r_[np.ones(m), -np.ones(n)])
+        V10 = perm(np.r_[np.ones(m), np.zeros(n)])
+        V01 = perm(np.r_[np.zeros(m), -np.ones(n)])
+
+    N = number_bandwidths * len(kbl)
+    if backend == "jax":
+        M = _jax_statistics(X, Y, kbl, wild, (Rn,) if wild else (V10, V01, V11))
+        bws = np.concatenate([b for _, b, _ in kbl])
+        names = [k for k, b, _ in kbl for _ in b]
+        kbl_loop = []
+    else:
+        M = np.empty((N, Btot)); bws = np.empty(N); names = []
+        kbl_loop = kbl
+    Z = np.vstack((X, Y))
+    last = None
+    for j, (kern, bandwidths, metric) in enumerate(kbl_loop):
+        if metric != last:
+            P = cdist(Z, Z, metric); last = metric
+        for i, bw in enumerate(bandwidths):
+            K = _kernel(P, kern, bw)
+            np.fill_diagonal(K, 0)
+            if wild:
+                np.fill_diagonal(K[:n, n:], 0); np.fill_diagonal(K[n:, :n], 0)
+                if fast:
+                    H = K[:n, :n] + K[n:, n:] - K[:n, n:] - K[n:, :n]
+                    M[j * number_bandwidths + i] = np.sum(Rn * (H @ Rn), 0)
+                else:
+                    M[j * number_bandwidths + i] = np.sum(R * (K @ R), 0)
+            else:
+                M[j * number_bandwidths + i] = (
+                    np.sum(V10 * (K @ V10), 0) * (n - m + 1) / (m * n * (m - 1))
+                    + np.sum(V01 * (K @ V01), 0) * (m - n + 1) / (m * n * (n - 1))
+                    + np.sum(V11 * (K @ V11), 0) / (m * n))
+            bws[j * number_bandwidths + i] = bw; names.append(kern)
+
+    stat = M[:, B1]
+    M1 = np.sort(M[:, :B1 + 1]); M2 = M[:, B1 + 1:]
+    allw = np.tile(w, len(kbl))
+    rows = np.arange(N)
+
+    def quantiles(u):
+        cols = np.ceil((B1 + 1) * (1 - u * allw)).astype(int) - 1
+        return M1[rows, cols]
+
+    u_lo, u_hi = 0.0, float(np.min(1 / w))
+    for _ in range(B3):
+        u = (u_lo + u_hi) / 2
+        if np.mean(np.max(M2 - quantiles(u)[:, None], 0) > 0) <= alpha:
+            u_lo = u
+        else:
+            u_hi = u
+    u = u_lo
+    q = quantiles(u)
+    reject_each = stat > q
+    out = int(reject_each.any())
+    if not return_details:
+        return out
+    pvals = np.mean(M1 - stat[:, None] >= 0, -1)
+    return out, {"kernel": names, "bandwidth": bws, "mmd": stat, "quantile": q,
+                 "p_value": pvals, "threshold": u * allw, "reject": reject_each}

kt/mmdfuse.py

@@ -0,0 +1,141 @@
+"""MMD-FUSE two-sample test (Biggs, Schrab, Gretton; NeurIPS 2023).
+
+NumPy re-implementation adapted from the MIT-licensed JAX reference package
+https://github.com/antoninschrab/mmdfuse (Copyright (c) 2023 Antonin Schrab).
+The statistic, bandwidth collection and permutation calibration follow the
+reference in the same evaluation order. JAX's random-number generator cannot be
+reproduced in NumPy, so the permutations are drawn from ``np.random.RandomState(seed)``;
+given the same permutations (private ``_permutations`` hook) the results agree with
+the reference up to floating-point rounding (see ``check_mmdfuse_vs_reference.py``).
+"""
+from __future__ import annotations
+import numpy as np
+from scipy.spatial.distance import cdist
+from scipy.special import logsumexp
+
+_L1 = ("laplace",) + tuple(f"matern_{i}.5_l1" for i in range(5))
+_L2 = ("imq", "rq", "gaussian") + tuple(f"matern_{i}.5_l2" for i in range(5))
+KERNELS = _L1 + _L2
+
+
+def _kernel(pairwise: np.ndarray, kernel: str, bandwidth: float, rq_exponent: float = 0.5) -> np.ndarray:
+    d = pairwise / bandwidth
+    if kernel == "gaussian":
+        return np.exp(-(d ** 2) / 2)
+    if kernel == "laplace":
+        return np.exp(-d * np.sqrt(2))
+    if kernel == "rq":
+        return (1 + d ** 2 / (2 * rq_exponent)) ** (-rq_exponent)
+    if kernel == "imq":
+        return (1 + d ** 2) ** (-0.5)
+    base = kernel[:10]
+    if base == "matern_0.5":
+        return np.exp(-d)
+    if base == "matern_1.5":
+        return (1 + np.sqrt(3) * d) * np.exp(-np.sqrt(3) * d)
+    if base == "matern_2.5":
+        return (1 + np.sqrt(5) * d + 5 / 3 * d ** 2) * np.exp(-np.sqrt(5) * d)
+    if base == "matern_3.5":
+        return (1 + np.sqrt(7) * d + 2 * 7 / 5 * d ** 2 + 7 * np.sqrt(7) / 3 / 5 * d ** 3) * np.exp(-np.sqrt(7) * d)
+    return (1 + 3 * d + 3 * (6 ** 2) / 28 * d ** 2 + (6 ** 3) / 84 * d ** 3
+            + (6 ** 4) / 1680 * d ** 4) * np.exp(-3 * d)  # matern_4.5
+
+
+def _bandwidths(distances: np.ndarray, number: int) -> np.ndarray:
+    median = np.median(distances)
+    distances = distances + (distances == 0) * median
+    dd = np.sort(distances)
+    lam_min = dd[int(np.floor(len(dd) * 0.05))] / 2
+    lam_max = dd[int(np.floor(len(dd) * 0.95))] * 2
+    return np.linspace(lam_min, lam_max, number)
+
+
+def mmdfuse(X, Y, alpha=0.05, kernels=("laplace", "gaussian"), lambda_multiplier=1,
+            number_bandwidths=10, number_permutations=2000, seed=42, return_details=False,
+            _permutations=None):
+    """Return 1 if MMD-FUSE rejects H0: P_X = P_Y at level alpha, else 0.
+
+    MMD-FUSE fuses normalised MMD estimates over ``kernels`` x ``number_bandwidths``
+    bandwidths with a soft maximum (log-mean-exp, temperature ``lambda_multiplier``) and
+    calibrates the fused statistic with ``number_permutations`` permutations.
+
+    kernels: a kernel name or tuple of names from ``"laplace"``, ``"gaussian"``, ``"imq"``,
+        ``"rq"``, ``"matern_{0..4}.5_l1"``, ``"matern_{0..4}.5_l2"``.
+    seed: seeds ``np.random.RandomState`` for the permutations.
+    return_details: also return a dict with ``statistic``, ``threshold`` (reject iff
+        ``statistic > threshold``), ``p_value``, ``permuted_statistics``, ``kernel``,
+        ``bandwidth`` and the normalised single-kernel ``mmd`` values.
+    """
+    X = np.asarray(X, dtype=float); Y = np.asarray(Y, dtype=float)
+    if X.ndim != 2 or Y.ndim != 2 or X.shape[1] != Y.shape[1]:
+        raise ValueError("X and Y must be 2-D arrays with the same number of columns")
+    if Y.shape[0] > X.shape[0]:
+        X, Y = Y, X
+    m, n = X.shape[0], Y.shape[0]
+    if n < 2:
+        raise ValueError("need at least 2 samples in each of X and Y")
+    if not 0 < alpha < 1:
+        raise ValueError("alpha must be in (0, 1)")
+    if not lambda_multiplier > 0:
+        raise ValueError("lambda_multiplier must be positive")
+    if not (isinstance(number_bandwidths, int) and number_bandwidths > 1):
+        raise ValueError("number_bandwidths must be an int > 1")
+    if not (isinstance(number_permutations, int) and number_permutations > 0):
+        raise ValueError("number_permutations must be a positive int")
+    if isinstance(kernels, str):
+        kernels = (kernels,)
+    kernels = tuple(kernels)
+    if not kernels or any(k not in KERNELS for k in kernels):
+        raise ValueError(f"kernels must be a non-empty collection from {KERNELS}")
+
+    B = number_permutations
+    if _permutations is None:
+        idx = np.random.RandomState(seed).rand(B, m + n).argsort(axis=1)
+    else:
+        idx = np.asarray(_permutations)
+        if idx.shape != (B, m + n):
+            raise ValueError("_permutations must have shape (number_permutations, m + n)")
+
+    def perm(v):
+        V = np.empty((B + 1, m + n)); V[:B] = v[idx]; V[B] = v
+        return V.T
+
+    V11 = perm(np.r_[np.ones(m), -np.ones(n)])
+    V10 = perm(np.r_[np.ones(m), np.zeros(n)])
+    V01 = perm(np.r_[np.zeros(m), -np.ones(n)])
+
+    # same ordering as the reference: all l1 kernels first, then l2
+    groups = [("cityblock", [k for k in kernels if k in _L1]),
+              ("euclidean", [k for k in kernels if k in _L2])]
+    N = number_bandwidths * len(kernels)
+    M = np.empty((N, B + 1)); bws = np.empty(N); names = []
+    Z = np.vstack((X, Y))
+    row = 0
+    for metric, ks in groups:
+        if not ks:
+            continue
+        P = cdist(Z, Z, metric)
+        bandwidths = _bandwidths(P[np.triu_indices(P.shape[0])], number_bandwidths)
+        for kern in ks:
+            for bw in bandwidths:
+                K = _kernel(P, kern, bw)
+                np.fill_diagonal(K, 0)
+                unscaled_std = np.sqrt(np.sum(K ** 2))
+                M[row] = ((np.sum(V10 * (K @ V10), 0) * (n - m + 1) * (n - 1) / (m * (m - 1))
+                           + np.sum(V01 * (K @ V01), 0) * (m - n + 1) / m
+                           + np.sum(V11 * (K @ V11), 0) * (n - 1) / m)
+                          / unscaled_std * np.sqrt(n * (n - 1)))
+                bws[row] = bw; names.append(kern); row += 1
+
+    all_stats = logsumexp(lambda_multiplier * M, axis=0, b=1 / N)
+    stat = all_stats[-1]
+    p_val = np.mean(all_stats >= stat)
+    out = int(p_val <= alpha)
+    if not return_details:
+        return out
+    # reject iff p_val <= alpha  <=>  stat > the ceil((B+1)(1-alpha))-th smallest statistic
+    k = int(np.ceil((B + 1) * (1 - alpha))) - 1
+    threshold = np.sort(all_stats)[k]
+    return out, {"statistic": stat, "threshold": threshold, "p_value": p_val,
+                 "permuted_statistics": all_stats[:-1], "kernel": names, "bandwidth": bws,
+                 "mmd": M[:, -1]}