edmkit 0.0.16__tar.gz → 0.0.18__tar.gz

edmkit-0.0.18/PKG-INFO

@@ -0,0 +1,87 @@
+Metadata-Version: 2.3
+Name: edmkit
+Version: 0.0.18
+Summary: Simple EDM (Empirical Dynamic Modeling) library
+Author: FUJISHIGE TEMMA
+Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>
+Requires-Dist: kdtree-rs>=0.3.1
+Requires-Dist: numpy>=2.5.4
+Requires-Dist: scipy>=1.18.1
+Requires-Dist: tinygrad>=0.14.0
+Requires-Dist: tinysolve>=0.3.0
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+
+# edmkit
+
+This library is a collection of tools and utilities that are useful for Empirical Data Modeling (EDM) and related tasks. The library is designed to be fast and lightweight, and easy to use.
+
+[benchmarks](./benchmarks)
+
+![](./benchmarks/benchmarks.png)
+
+[compare.py](./benchmarks/compare.py)
+
+![](./benchmarks/compare.png)
+
+Measured on an Apple M4 Pro on November 4, 2026. Times are medians from `pytest-benchmark`. Parallel CCM uses 14 workers. Peak memory is process-tree RSS in MiB, sampled separately from timing for each case. It includes the interpreter, input, and worker processes. Summed RSS can count shared pages in more than one process.
+
+| simplex N | edmkit [ms] | pyEDM [ms] | speedup | edmkit peak [MiB] | pyEDM peak [MiB] |
+| ---: | ---: | ---: | ---: | ---: | ---: |
+| 500 | 0.17 | 1.5 | 9.0× | 168 | 169 |
+| 1,000 | 0.39 | 2.4 | 6.1× | 169 | 170 |
+| 2,000 | 0.81 | 4.4 | 5.5× | 168 | 170 |
+| 5,000 | 2.2 | 11.5 | 5.3× | 170 | 175 |
+| 10,000 | 4.7 | 23.8 | 5.0× | 172 | 180 |
+| 20,000 | 10.2 | 50.3 | 4.9× | 176 | 193 |
+
+| smap N | edmkit [ms] | pyEDM [ms] | speedup | edmkit peak [MiB] | pyEDM peak [MiB] |
+| ---: | ---: | ---: | ---: | ---: | ---: |
+| 500 | 2.1 | 36.6 | 17.9× | 176 | 180 |
+| 1,000 | 8.7 | 128.3 | 14.7× | 192 | 216 |
+| 2,000 | 22.0 | 513.6 | 23.4× | 261 | 357 |
+| 5,000 | 193.9 | 3,376.7 | 17.4× | 744 | 1,344 |
+
+| ccm N | edmkit [ms] | pyEDM [ms] | speedup | edmkit parallel [ms] | pyEDM parallel [ms] | speedup |
+| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| 500 | 106.4 | 151.7 | 1.4× | 19.5 | 68.0 | 3.5× |
+| 1,000 | 168.4 | 327.0 | 1.9× | 26.3 | 91.9 | 3.5× |
+| 2,000 | 329.7 | 745.1 | 2.3× | 48.8 | 148.3 | 3.0× |
+| 5,000 | 668.9 | 1,945.7 | 2.9× | 97.4 | 397.5 | 4.1× |
+| 10,000 | 1,253.6 | 4,176.7 | 3.3× | 164.6 | 585.4 | 3.6× |
+
+| ccm N | edmkit peak [MiB] | pyEDM peak [MiB] | edmkit parallel peak [MiB] | pyEDM parallel peak [MiB] |
+| ---: | ---: | ---: | ---: | ---: |
+| 500 | 171 | 170 | 191 | 1,417 |
+| 1,000 | 173 | 176 | 196 | 1,429 |
+| 2,000 | 179 | 186 | 207 | 1,460 |
+| 5,000 | 204 | 224 | 239 | 1,624 |
+| 10,000 | 230 | 362 | 297 | 1,804 |
+
+
+## Installation
+
+To install the library, you can use pip:
+
+```bash
+pip install edmkit
+```
+
+Or you can also use [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv add edmkit
+```
+
+## Testing
+
+```bash
+# Run all tests
+uv run pytest
+
+# Skip slow tests
+uv run pytest -m "not slow"
+
+# Run with CI profile (more hypothesis examples)
+HYPOTHESIS_PROFILE=ci uv run pytest
+```

edmkit-0.0.18/README.md

@@ -0,0 +1,73 @@
+# edmkit
+
+This library is a collection of tools and utilities that are useful for Empirical Data Modeling (EDM) and related tasks. The library is designed to be fast and lightweight, and easy to use.
+
+[benchmarks](./benchmarks)
+
+![](./benchmarks/benchmarks.png)
+
+[compare.py](./benchmarks/compare.py)
+
+![](./benchmarks/compare.png)
+
+Measured on an Apple M4 Pro on November 4, 2026. Times are medians from `pytest-benchmark`. Parallel CCM uses 14 workers. Peak memory is process-tree RSS in MiB, sampled separately from timing for each case. It includes the interpreter, input, and worker processes. Summed RSS can count shared pages in more than one process.
+
+| simplex N | edmkit [ms] | pyEDM [ms] | speedup | edmkit peak [MiB] | pyEDM peak [MiB] |
+| ---: | ---: | ---: | ---: | ---: | ---: |
+| 500 | 0.17 | 1.5 | 9.0× | 168 | 169 |
+| 1,000 | 0.39 | 2.4 | 6.1× | 169 | 170 |
+| 2,000 | 0.81 | 4.4 | 5.5× | 168 | 170 |
+| 5,000 | 2.2 | 11.5 | 5.3× | 170 | 175 |
+| 10,000 | 4.7 | 23.8 | 5.0× | 172 | 180 |
+| 20,000 | 10.2 | 50.3 | 4.9× | 176 | 193 |
+
+| smap N | edmkit [ms] | pyEDM [ms] | speedup | edmkit peak [MiB] | pyEDM peak [MiB] |
+| ---: | ---: | ---: | ---: | ---: | ---: |
+| 500 | 2.1 | 36.6 | 17.9× | 176 | 180 |
+| 1,000 | 8.7 | 128.3 | 14.7× | 192 | 216 |
+| 2,000 | 22.0 | 513.6 | 23.4× | 261 | 357 |
+| 5,000 | 193.9 | 3,376.7 | 17.4× | 744 | 1,344 |
+
+| ccm N | edmkit [ms] | pyEDM [ms] | speedup | edmkit parallel [ms] | pyEDM parallel [ms] | speedup |
+| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| 500 | 106.4 | 151.7 | 1.4× | 19.5 | 68.0 | 3.5× |
+| 1,000 | 168.4 | 327.0 | 1.9× | 26.3 | 91.9 | 3.5× |
+| 2,000 | 329.7 | 745.1 | 2.3× | 48.8 | 148.3 | 3.0× |
+| 5,000 | 668.9 | 1,945.7 | 2.9× | 97.4 | 397.5 | 4.1× |
+| 10,000 | 1,253.6 | 4,176.7 | 3.3× | 164.6 | 585.4 | 3.6× |
+
+| ccm N | edmkit peak [MiB] | pyEDM peak [MiB] | edmkit parallel peak [MiB] | pyEDM parallel peak [MiB] |
+| ---: | ---: | ---: | ---: | ---: |
+| 500 | 171 | 170 | 191 | 1,417 |
+| 1,000 | 173 | 176 | 196 | 1,429 |
+| 2,000 | 179 | 186 | 207 | 1,460 |
+| 5,000 | 204 | 224 | 239 | 1,624 |
+| 10,000 | 230 | 362 | 297 | 1,804 |
+
+
+## Installation
+
+To install the library, you can use pip:
+
+```bash
+pip install edmkit
+```
+
+Or you can also use [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv add edmkit
+```
+
+## Testing
+
+```bash
+# Run all tests
+uv run pytest
+
+# Skip slow tests
+uv run pytest -m "not slow"
+
+# Run with CI profile (more hypothesis examples)
+HYPOTHESIS_PROFILE=ci uv run pytest
+```

{edmkit-0.0.16 → edmkit-0.0.18}/pyproject.toml

@@ -1,15 +1,15 @@
 [project]
 name = "edmkit"
-version = "0.0.16"
+version = "0.0.18"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 readme = "README.md"
 requires-python = ">= 3.13"
 dependencies = [
     "kdtree-rs>=0.3.1",
-    "numpy>=2.5.1",
-    "scipy>=1.17.1",
+    "numpy>=2.5.4",
+    "scipy>=1.18.1",
     "tinygrad>=0.14.0",
-    "tinysolve>=0.2.1",
+    "tinysolve>=0.3.0",
 ]
 
 [[project.authors]]
@@ -17,14 +17,19 @@ name = "FUJISHIGE TEMMA"
 email = "tenma.x0@gmail.com"
 
 [dependency-groups]
+benchmark = [
+    "pandas>=3.1.0",
+    "pyedm>=2.5.7",
+]
 dev = [
-    "griffe2md>=1.5.0",
-    "hypothesis>=6.151.9",
+    "griffe2md>=1.5.1",
+    "hypothesis>=6.170.14",
     "matplotlib>=3.11.2",
-    "pytest>=9.1.1",
+    "pytest>=9.2.0",
+    "pytest-benchmark>=5.3.1",
     "pywavelets>=1.10.0",
-    "ruff>=0.15.20",
-    "ty>=0.0.56",
+    "ruff>=0.17.3",
+    "ty>=0.0.95",
 ]
 
 [tool.uv]
@@ -36,10 +41,17 @@ namespace = true
 
 [tool.pytest.ini_options]
 addopts = "-v --tb=short"
+testpaths = ["tests"]
+python_files = [
+    "test_*.py",
+    "bench_*.py",
+]
 filterwarnings = ["error"]
 markers = [
     """slow: marks tests as slow (deselect with '-m "not slow"')""",
     """gpu: marks tests requiring tinygrad tensor backend (deselect with '-m "not gpu"')""",
+    "dense: marks benchmarks that build an (M, N) matrix",
+    """huge: marks benchmarks whose N is too large for an (M, N) matrix in memory (deselect with '-m "not (dense and huge)"')""",
 ]
 
 [build-system]

{edmkit-0.0.16 → edmkit-0.0.18}/pyproject.toml.orig

@@ -1,27 +1,32 @@
 [project]
 name = "edmkit"
-version = "0.0.16"
+version = "0.0.18"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"
 requires-python = ">= 3.13"
 dependencies = [
     "kdtree-rs>=0.3.1",
-    "numpy>=2.5.1",
-    "scipy>=1.17.1",
+    "numpy>=2.5.4",
+    "scipy>=1.18.1",
     "tinygrad>=0.14.0",
-    "tinysolve>=0.2.1",
+    "tinysolve>=0.3.0",
 ]
 
 [dependency-groups]
+benchmark = [
+    "pandas>=3.1.0",
+    "pyedm>=2.5.7",
+]
 dev = [
-    "griffe2md>=1.5.0",
-    "hypothesis>=6.151.9",
+    "griffe2md>=1.5.1",
+    "hypothesis>=6.170.14",
     "matplotlib>=3.11.2",
-    "pytest>=9.1.1",
+    "pytest>=9.2.0",
+    "pytest-benchmark>=5.3.1",
     "pywavelets>=1.10.0",
-    "ruff>=0.15.20",
-    "ty>=0.0.56",
+    "ruff>=0.17.3",
+    "ty>=0.0.95",
 ]
 
 [tool.uv]
@@ -31,10 +36,14 @@ required-environments = [
 
 [tool.pytest.ini_options]
 addopts = "-v --tb=short"
+testpaths = ["tests"]
+python_files = ["test_*.py", "bench_*.py"]
 filterwarnings = ["error"]
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",
     "gpu: marks tests requiring tinygrad tensor backend (deselect with '-m \"not gpu\"')",
+    "dense: marks benchmarks that build an (M, N) matrix",
+    "huge: marks benchmarks whose N is too large for an (M, N) matrix in memory (deselect with '-m \"not (dense and huge)\"')",
 ]
 
 

edmkit-0.0.18/src/edmkit/ccm.py

@@ -0,0 +1,168 @@
+from collections.abc import Callable, Iterator
+from functools import partial
+
+import numpy as np
+
+from edmkit.types import Mask, PredictFunc
+
+__all__ = ["bootstrap"]
+
+
+def bootstrap(
+    X: np.ndarray,
+    Y: np.ndarray,
+    Q: np.ndarray,
+    lib_sizes: np.ndarray,
+    *,
+    predict_func: PredictFunc,
+    n_samples: int = 20,
+    replace: bool = True,
+    seed: int | None = 42,
+    mask: Mask | None = None,
+    batch_size: int | None = 20,
+) -> Iterator[tuple[tuple[int, slice], Callable[[], np.ndarray]]]:
+    """
+    Iterate `predict_func(X[library], Y[library], Q)` for `n_samples` times for each library size in `lib_sizes`.
+
+    By gathering the results of the returned tasks, you can get `(len(lib_sizes), n_samples, M, E')` array of predictions.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        The input data of shape (N,) or (N, E)
+    Y : np.ndarray
+        The target data of shape (N,) or (N, E'), sampled together with `X`.
+    Q : np.ndarray
+        The query points of shape (M,) or (M, E), shared by all sampled libraries as a read-only view.
+    lib_sizes : np.ndarray
+        The 1D array of library sizes, yielded in the given order. Sizes may exceed N only if `replace` is True.
+    predict_func : PredictFunc
+        The batched prediction function, such as `simplex_projection` or ``functools.partial(smap, theta=...)``.
+    n_samples : int, default 20
+        The number of libraries to sample for each size in `lib_sizes`.
+    replace : bool, default True
+        Whether to sample the library points with replacement.
+    seed : int or None, default 42
+        The seed of the random generator. If None, uses fresh entropy.
+    mask : Mask or None
+        ``mask(i, j)`` indicating, for each query in `Q`, which points of `X` to include. Use `theiler_window` for temporal exclusion.
+    batch_size : int or None, default 20
+        The maximum number of libraries per task. If None, uses `n_samples`.
+
+    Yields
+    ------
+    index : tuple[int, slice]
+        The position of the task's libraries in an array of shape (len(`lib_sizes`), `n_samples`), such as the scores.
+    samples : Callable[[], np.ndarray]
+        The task that returns the predictions on B sampled libraries, of shape (B, M, E'), where B <= `batch_size` and E' = 1 if `Y` is 1D.
+
+    Raises
+    ------
+    ValueError
+        - If `X`, `Y`, or `Q` is empty or not 1D or 2D.
+        - If the input arrays `X` and `Y` do not have the same number of points.
+        - If the dimension of `X` and `Q` do not match.
+        - If `lib_sizes` is not a nonempty 1D array of positive integers.
+        - If some size in `lib_sizes` exceeds N when `replace` is False.
+        - If `n_samples` or `batch_size` is not positive.
+        - If `predict_func` raises when a task is executed.
+
+    Examples
+    --------
+    ```python
+    from concurrent.futures import ThreadPoolExecutor
+
+    import numpy as np
+
+    from edmkit.ccm import bootstrap
+    from edmkit.embedding import embed
+    from edmkit.metrics import pearson_correlation
+    from edmkit.simplex_projection import simplex_projection
+
+    # Generate a simple time series (logistic map)
+    N = 300
+    x = np.zeros(N)
+    x[0] = 0.4
+    for i in range(1, N):
+        x[i] = 3.9 * x[i - 1] * (1 - x[i - 1])
+
+    tau = 2
+    E = 3
+
+    coordinates = np.array([[0, tau * j] for j in range(E)])
+    embedding, _ = embed(coordinates, x)
+    shift = tau * (E - 1)
+
+    pool_size = 200
+    Tp = 1
+    X = embedding[:pool_size - shift]
+    Y = x[shift + Tp : pool_size + Tp]
+    Q = embedding[pool_size - shift : -Tp]
+    actual = x[pool_size + Tp :, None]
+
+    lib_sizes = np.array([10, 50, 190])
+    n_samples = 100
+
+
+    def score(samples):
+        predictions = samples()  # (B, M, 1)
+        return pearson_correlation(predictions, np.broadcast_to(actual, predictions.shape))
+
+
+    scores = np.empty((len(lib_sizes), n_samples))
+    tasks = bootstrap(X, Y, Q, lib_sizes, predict_func=simplex_projection, n_samples=n_samples)
+    indices, samples = zip(*tasks)
+    with ThreadPoolExecutor() as executor:
+        for index, result in zip(indices, executor.map(score, samples)):
+            scores[index] = result
+
+    for lib_size, correlation in zip(lib_sizes, scores.mean(axis=1)):
+        print(f"Correlation (L={lib_size}): {correlation:.3f}")
+    ```
+    """
+    if X.ndim == 1:
+        X = X[:, None]
+    if Y.ndim == 1:
+        Y = Y[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if X.ndim != 2 or Y.ndim != 2 or Q.ndim != 2:
+        raise ValueError(f"X, Y, and Q must be 1D or 2D arrays, got X.ndim={X.ndim}, Y.ndim={Y.ndim}, Q.ndim={Q.ndim}")
+    if X.size == 0 or Y.size == 0 or Q.size == 0:
+        raise ValueError("X, Y, and Q must be nonempty")
+    if X.shape[0] != Y.shape[0]:
+        raise ValueError(f"X and Y must have the same length, got X.shape={X.shape}, Y.shape={Y.shape}")
+    if X.shape[1] != Q.shape[1]:
+        raise ValueError(f"X and Q must have the same dimension, got X.shape={X.shape}, Q.shape={Q.shape}")
+    if lib_sizes.ndim != 1 or lib_sizes.size == 0 or not np.issubdtype(lib_sizes.dtype, np.integer) or np.any(lib_sizes <= 0):
+        raise ValueError("lib_sizes must be a nonempty 1D integer array of positive sizes")
+    if not replace and np.any(lib_sizes > len(X)):
+        raise ValueError(f"lib_sizes cannot exceed N={len(X)} when replace is False")
+    batch_size = n_samples if batch_size is None else batch_size
+    if n_samples <= 0 or batch_size <= 0:
+        raise ValueError(f"n_samples and batch_size must be positive, got n_samples={n_samples}, batch_size={batch_size}")
+
+    rng = np.random.default_rng(seed)
+    for i, lib_size in enumerate(lib_sizes):
+        for start in range(0, n_samples, batch_size):
+            batch = min(batch_size, n_samples - start)
+            if replace:
+                indices = rng.choice(len(X), size=(batch, lib_size))
+            else:
+                indices = np.stack([rng.choice(len(X), size=lib_size, replace=False) for _ in range(batch)])
+            yield (i, slice(start, start + batch)), partial(samples, X, Y, Q, indices, predict_func=predict_func, mask=mask)
+
+
+def samples(
+    X: np.ndarray,
+    Y: np.ndarray,
+    Q: np.ndarray,
+    indices: np.ndarray,
+    *,
+    predict_func: PredictFunc,
+    mask: Mask | None,
+) -> np.ndarray:
+    queries = np.broadcast_to(Q, (len(indices), *Q.shape))
+    # position j of library b is the point indices[b, j] of X
+    return predict_func(X[indices], Y[indices], queries, mask=None if mask is None else (lambda b, i, j: mask(i, indices[b, j])))

edmkit-0.0.18/src/edmkit/simplex_projection/knn.py

@@ -0,0 +1,65 @@
+import numpy as np
+from kdtree import KDTree
+
+from edmkit.types import Mask
+
+__all__: list[str] = []
+
+
+def knn(X: np.ndarray, Q: np.ndarray, k: int, mask: Mask | None = None) -> tuple[np.ndarray, np.ndarray]:
+    """Find the k-nearest neighbors of `Q` in `X` using `kdtree.KDTree`, optionally restricted by `mask`.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        The input data (N, E)
+    Q : np.ndarray
+        The query points (M, E)
+    k : int
+        The number of nearest neighbors to find (typically E+1 for simplex projection).
+    mask : Mask or None
+        ``mask(i, j)`` is True where X[j] is a valid point for Q[i].
+
+    Returns
+    -------
+    distances : np.ndarray
+        The distances from each query point in `Q` to its k nearest neighbors in `X`, sorted in ascending order (M, k)
+    indices : np.ndarray
+        The indices of the k nearest neighbors in `X` for each query point in `Q` (M, k)
+    """
+    N = X.shape[0]
+
+    if N < k:
+        raise ValueError(f"Not enough points in X to find {k} neighbors, got N={N}")
+
+    # smaller leaves improve pruning in low dimensions.
+    tree = KDTree(X, leafsize=16 if X.shape[-1] <= 8 else 32)
+    if mask is None or len(Q) == 0:
+        return tree.query(Q, k=k)
+
+    distances = np.empty((Q.shape[0], k))
+    indices = np.empty((Q.shape[0], k), dtype=np.intp)
+
+    pending = np.arange(Q.shape[0])
+    k_query = min(k + 1, N)
+
+    while pending.size > 0:
+        _distances, _indices = tree.query(Q if pending.size == len(Q) else Q[pending], k=k_query)
+
+        # evaluate the mask and keep the first k valid candidates.
+        valid = np.broadcast_to(np.asarray(mask(pending[:, None], _indices), dtype=bool), _indices.shape)
+        count = valid.sum(axis=-1)
+
+        if k_query == N and count.min() < k:
+            raise ValueError(f"Not enough unmasked points in X to find {k} neighbors, got {count.min()} for some query")
+
+        done = count >= k
+        if done.any():
+            first_k = valid & (valid.cumsum(axis=-1, dtype=np.uint32) <= k) & done[:, None]
+            distances[pending[done]] = _distances[first_k].reshape(-1, k)
+            indices[pending[done]] = _indices[first_k].reshape(-1, k)
+
+        pending = pending[~done]
+        k_query = min(2 * k_query, N)
+
+    return distances, indices

{edmkit-0.0.16 → edmkit-0.0.18}/src/edmkit/simplex_projection/linear_operator.py

@@ -1,17 +1,20 @@
 from __future__ import annotations
 
+from functools import partial
 from typing import TYPE_CHECKING, overload
 
 import numpy as np
 
 from edmkit.simplex_projection.knn import knn
-from edmkit.util import pairwise_distance
+from edmkit.util import materialize, pairwise_distance, topk
 
 __all__ = ["linear_operator"]
 
 if TYPE_CHECKING:
     from tinygrad import Tensor
 
+    from edmkit.types import Mask
+
 
 @overload
 def linear_operator(
@@ -19,7 +22,7 @@ def linear_operator(
     Q: np.ndarray,
     *,
     k: int | None = None,
-    mask: np.ndarray | None = None,
+    mask: Mask | None = None,
 ) -> np.ndarray: ...
 
 
@@ -29,7 +32,7 @@ def linear_operator(
     Q: Tensor,
     *,
     k: int | None = None,
-    mask: Tensor | None = None,
+    mask: Mask | None = None,
 ) -> Tensor: ...
 
 
@@ -51,8 +54,8 @@ def linear_operator(
         The query points of shape (M,) or (M, E) or (B, M, E) for which to find the nearest neighbors in `X`.
     k : int or None, default None
         The number of nearest neighbors to use. If None, uses E + 1, where E is the dimension of `X`.
-    mask : np.ndarray or Tensor or None
-        Boolean mask of shape (M, N) or (B, M, N) indicating, for each query in `Q`, which library points to include when finding nearest neighbors.
+    mask : Mask or None
+        ``mask(i, j)``, or ``mask(b, i, j)``, indicating for each query in `Q` which library points to include when finding nearest neighbors.
 
     Returns
     -------
@@ -64,7 +67,6 @@ def linear_operator(
     ValueError
         - If `k` is not positive.
         - If the batch size or dimension of `X` and `Q` do not match.
-        - If `mask` does not have shape (M, N) or (B, M, N).
         - If fewer than `k` library points are unmasked for some query.
 
     Examples
@@ -93,7 +95,7 @@ def _numpy(
     Q: np.ndarray,
     *,
     k: int | None = None,
-    mask: np.ndarray | None = None,
+    mask: Mask | None = None,
 ) -> np.ndarray:
     # ensure at least 2D
     if X.ndim == 1:
@@ -105,8 +107,6 @@ def _numpy(
         raise ValueError(f"X and Q must both be 2D or both be 3D arrays, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
     if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
         raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
-    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
-        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
 
     return_ndim = X.ndim
 
@@ -114,7 +114,7 @@ def _numpy(
     if X.ndim == 2:
         X, Q = X[None], Q[None]
         if mask is not None:
-            mask = mask[None]
+            mask = lambda _, *index, mask=mask: mask(*index)
 
     B, N, E = X.shape
     M = Q.shape[1]
@@ -125,19 +125,7 @@ def _numpy(
     distances = np.empty((B, M, k))
     indices = np.empty((B, M, k), dtype=np.intp)
     for b in range(B):
-        if mask is None:
-            distances[b], indices[b] = knn(X[b], Q[b], k)
-        else:  # over-fetch, then drop the masked-out candidates
-            n_exclude = int((~mask[b]).sum(axis=-1).max())
-            if N - n_exclude < k:
-                raise ValueError(
-                    f"Not enough unmasked points in X to find {k} neighbors, got N={N} with up to {n_exclude} points masked out per query"
-                )
-            _distances, _indices = knn(X[b], Q[b], k + n_exclude)
-            _distances = np.where(np.take_along_axis(mask[b], _indices, axis=-1), _distances, np.inf)
-            top_k = np.argsort(_distances, axis=-1)[..., :k]
-            distances[b] = np.take_along_axis(_distances, top_k, axis=-1)
-            indices[b] = np.take_along_axis(_indices, top_k, axis=-1)
+        distances[b], indices[b] = knn(X[b], Q[b], k, None if mask is None else partial(mask, b))
 
     # clamp to avoid division by zero
     d_min = np.fmax(distances.min(axis=-1, keepdims=True), 1e-6)  # (B, M, 1)
@@ -157,8 +145,10 @@ def _tensor(
     Q: Tensor,
     *,
     k: int | None = None,
-    mask: Tensor | None = None,
+    mask: Mask | None = None,
 ) -> Tensor:
+    from tinygrad import Tensor
+
     # ensure at least 2D
     if X.ndim == 1:
         X = X[:, None]
@@ -169,16 +159,15 @@ def _tensor(
         raise ValueError(f"X and Q must both be 2D or both be 3D tensors, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
     if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
         raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
-    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
-        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
+    _mask = None if mask is None else Tensor(materialize(mask, Q.shape[:-1] + X.shape[-2:-1]))
 
     return_ndim = X.ndim
 
     # treat the unbatched case as a single batch
     if X.ndim == 2:
         X, Q = X[None], Q[None]
-        if mask is not None:
-            mask = mask[None]
+        if _mask is not None:
+            _mask = _mask[None]
 
     B, N, E = (int(X.shape[0]), int(X.shape[1]), int(X.shape[2]))
     M = int(Q.shape[1])
@@ -189,10 +178,10 @@ def _tensor(
     # clamp to avoid NaN gradient
     D = pairwise_distance(Q, X).clamp(min_=1e-12).sqrt()  # (B, M, N)
 
-    if mask is not None:
-        D = mask.where(D, float("inf"))
+    if _mask is not None:
+        D = _mask.where(D, float("inf"))
 
-    distances, indices = D.topk(k, dim=-1, largest=False, sorted_=True)  # (B, M, k)
+    distances, indices = topk(D, k, largest=False)  # (B, M, k)
 
     # neighbors are sorted, so the first is the nearest; clamp to avoid division by zero
     d_min = distances[..., :1].clamp(min_=1e-6)  # (B, M, 1)