[SimpingOjou]

Summary

This PR introduces a memory-efficient and faster chunked nearest-neighbor retrieval method (argmin) to DistanceMatrix, and updates ContinuousIndex.search to use it.

When training continuous supervised CEBRA (conditional='time_delta' or conditional='delta_normal') on large datasets (e.g., $N > 10^6$ samples), CEBRA computes pairwise distances between all candidate dataset points ($N$) and batch query points ($B$). In the current implementation, DistanceMatrix.__call__ allocates an entire $(N, B)$ float32 tensor at once. For large-scale continuous recordings (e.g., $N = 7.67 \times 10^6$ points with $B = 512$), this single matrix requires $\approx 14.6\text{ GiB}$ of contiguous VRAM, causing out-of-memory crashes (torch.OutOfMemoryError) on 16–24 GB GPUs.

By chunking candidate points along $N$ (default: 500,000 samples) and maintaining running minimums:

1. Peak VRAM footprint drops from $> 22\text{ GB}$ to $\approx 3.96\text{ GB}$, allowing continuous datasets to train comfortably on standard GPUs.
2. Pairwise computation replaces torch.einsum with optimized in-place BLAS GEMM (torch.mm), speeding up CPU/GPU search.

---

Changes

• cebra/distributions/index.py:
  • Added chunk_size: int = 500_000 parameter to DistanceMatrix.__init__.
  • Added DistanceMatrix.argmin(self, query, mask=None) -> torch.Tensor performing chunk-wise running minimums without materializing the full $(N, B)$ tensor in memory.
  • Updated ContinuousIndex.search(self, query) to delegate directly to self.dist_matrix.argmin(query).

---

Test Results

Tested on simulated data with fixed seeds with CPU:

--- Training Chunked CEBRA (Synchronized) ---
Unchunked Final Loss: 2.531332. Time: 82.395 s
Chunked Final Loss:   2.532491. Time: 66.460 s
Speedup: 1.24x

Ran the distribution test suite:

pytest tests/test_distributions.py
# ======================== 26 passed, 4 skipped in 6.04s =========================