From 6cf65ec9223bbc10173d9911fe4d65fea079273f Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 1 Aug 2026 02:33:41 +0000
Subject: [PATCH 01/12] Initial plan


From ef34f50d67000783cd351cdce7dfed14e7028291 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 1 Aug 2026 02:52:52 +0000
Subject: [PATCH 02/12] feat: replace chunk_iter with automatic iteration count
 derived from chunk_size

- Remove `chunk_iter` parameter from `pted()`, `pted_coverage_test()`,
  `pted_chunk_numpy/torch/jax`, and `_energy_distance_estimate_*` helpers
- Compute number of iterations automatically as `max(len(x), len(y)) // chunk_size`
- Iterate sequentially over the larger dataset; cycle through the smaller
  one using modular indexing when sizes differ
- Add `chunk_size <= max(len(x), len(y))` assertion in `pted()`
- Update all docstrings and README to reflect the new interface
- Update tests: remove `chunk_iter` args, fix seeds, add
  `test_pted_chunk_mismatched_sizes` for cycling behavior
---
 README.md          | 26 +++++++++---------
 src/pted/pted.py   | 59 ++++++++++++++++++-----------------------
 src/pted/utils.py  | 66 +++++++++++++++++++++++++---------------------
 tests/test_pted.py | 51 ++++++++++++++++++++++-------------
 4 files changed, 106 insertions(+), 96 deletions(-)

diff --git a/README.md b/README.md
index f460e36..9b6c033 100644
--- a/README.md
+++ b/README.md
@@ -275,7 +275,6 @@ def pted(
     metric: Union[str, float] = "euclidean",
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     two_tailed: bool = True,
     prog_bar: bool = False,
 ) -> Union[float, tuple[float, np.ndarray, float]]:
@@ -286,8 +285,7 @@ def pted(
 * **permutations** *(int)*: number of permutations to run. This determines how accurately the p-value is computed.
 * **metric** *(Union[str, float])*: distance metric to use. See scipy.spatial.distance.cdist for the list of available metrics with numpy. See torch.cdist when using PyTorch, note that the metric is passed as the "p" for torch.cdist and therefore is a float from 0 to inf. When using JAX arrays, the metric is passed as the "ord" for jnp.linalg.norm and therefore is also a float.
 * **return_all** *(bool)*: if True, return the test statistic and the permuted statistics with the p-value. If False, just return the p-value. bool (default: False)
-* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. This is useful for large datasets. The chunk size is the number of samples to use for each chunk. If None, use the full dataset.
-* **chunk_iter** *(Optional[int])*: The chunk iter is the number of iterations to use with the given chunk size.
+* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. The chunk size is the number of samples per chunk. The number of chunks is determined automatically as `max(len(x), len(y)) // chunk_size`, iterating over the larger dataset once and cycling through the smaller one. If None, use the full dataset.
 * **two_tailed** *(bool)*: if True, compute a two-tailed p-value. This is useful if you want to reject the null hypothesis when x and y are either too similar or too different. If False, only checks for dissimilarity but is more sensitive. Default is True.
 * **prog_bar** *(bool)*: if True, show a progress bar to track the progress of permutation tests. Default is False.
 
@@ -302,7 +300,6 @@ def pted_coverage_test(
     warn_confidence: Optional[float] = 1e-3,
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     sbc_histogram: Optional[str] = None,
     sbc_bins: Optional[int] = None,
     pit_plot: Optional[str] = None,
@@ -316,8 +313,7 @@ def pted_coverage_test(
 * **permutations** *(int)*: number of permutations to run. This determines how accurately the p-value is computed.
 * **metric** *(Union[str, float])*: distance metric to use. See scipy.spatial.distance.cdist for the list of available metrics with numpy. See torch.cdist when using PyTorch, note that the metric is passed as the "p" for torch.cdist and therefore is a float from 0 to inf. When using JAX arrays, the metric is passed as the "ord" for jnp.linalg.norm and therefore is also a float.
 * **return_all** *(bool)*: if True, return the test statistic and the permuted statistics with the p-value. If False, just return the p-value. bool (default: False)
-* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. This is useful for large datasets. The chunk size is the number of samples to use for each chunk. If None, use the full dataset.
-* **chunk_iter** *(Optional[int])*: The chunk iter is the number of iterations to use with the given chunk size.
+* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. The chunk size is the number of samples per chunk. The number of chunks is determined automatically as `max(len(x), len(y)) // chunk_size`, iterating over the larger dataset once and cycling through the smaller one. If None, use the full dataset.
 * **sbc_histogram** *(Optional[str])*: If given, the path/filename to save a Simulation-Based-Calibration histogram.
 * **sbc_bins** *(Optional[int])*: If given, force the histogram to have the provided number of bins. Otherwise, select an appropriate size: ~sqrt(N).
 * **pit_plot** *(Optional[str])*: If given, the path/filename to save a Probability Integral Transform (PIT) plot of the per-simulation p-values against the expected uniform distribution, with a shaded KS confidence band.
@@ -361,16 +357,18 @@ If a GPU isn't enough to get PTED running fast enough for you, or if you are
 running into memory limitations, there are still options! We can use an
 approximation of the energy distance, in this case the test is still exact but
 less sensitive than it would be otherwise. We can approximate the energy
-distance by taking random subsamples (chunks) of the full dataset, computing the
-energy distance, then averaging. Just set the `chunk_size` parameter for the
-number of samples you can manage at once and set the `chunk_iter` for the number
-of trials you want in the average. The larger these numbers are, the closer the
-estimate will be to the true energy distance, but it will take more compute.
-This lets you decide how to trade off compute for sensitivity.
+distance by iterating through sequential chunks of the full dataset, computing
+the energy distance on each chunk, then averaging. Just set the `chunk_size`
+parameter for the number of samples you can manage at once. The number of
+iterations is determined automatically as `max(len(x), len(y)) // chunk_size`,
+iterating over the larger dataset once and cycling through the smaller one if
+their sizes differ. The larger the chunk size, the closer the estimate will be
+to the true energy distance, but it will take more compute.
 
 Note that the computational complexity for standard PTED goes as 
 `O((n_samp_x + n_samp_y)^2)` while the chunked version goes as 
-`O(chunk_iter * (2 * chunk_size)^2)` so plan your chunking accordingly.
+`O(n_iter * (2 * chunk_size)^2)` where `n_iter = max(n_samp_x, n_samp_y) // chunk_size`,
+so plan your chunking accordingly.
 
 Example:
 ```python
@@ -380,7 +378,7 @@ import numpy as np
 x = np.random.normal(size = (500, 10)) # (n_samples_x, n_dimensions)
 y = np.random.normal(size = (400, 10)) # (n_samples_y, n_dimensions)
 
-p_value = pted(x, y, chunk_size = 50, chunk_iter = 100)
+p_value = pted(x, y, chunk_size = 50)
 print(f"p-value: {p_value:.3f}") # expect uniform random from 0-1
 ```
 
diff --git a/src/pted/pted.py b/src/pted/pted.py
index 0daf3b8..f30aa58 100644
--- a/src/pted/pted.py
+++ b/src/pted/pted.py
@@ -27,7 +27,6 @@ def pted(
     metric: Union[str, float] = "euclidean",
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     two_tailed: bool = True,
     prog_bar: bool = False,
 ) -> Union[float, tuple[float, np.ndarray, float]]:
@@ -92,10 +91,10 @@ def pted(
             bool (default: False)
         chunk_size (Optional[int]): if not None, use chunked energy distance
             estimation. This is useful for large datasets. The chunk size is the
-            number of samples to use for each chunk. If None, use the full
-            dataset.
-        chunk_iter (Optional[int]): The chunk iter is the number of iterations
-            to use with the given chunk size.
+            number of samples per chunk. The number of chunks is determined
+            automatically as ``max(len(x), len(y)) // chunk_size``, iterating
+            over the larger dataset once and cycling through the smaller one.
+            If None, use the full dataset.
         two_tailed (bool): if True, compute a two-tailed p-value. This is useful
             if you want to reject the null hypothesis when x and y are either
             too similar or too different. Default is True.
@@ -109,14 +108,13 @@ def pted(
         samples in x and y, D is the number of dimensions, and P is the number
         of permutations. For large datasets this can get unwieldy, so chunking
         is recommended. For chunking, the energy distance will be estimated at
-        each iteration rather than fully computed. To estimate the energy
-        distance, we take `chunk_size` sub-samples from x and y, and compute the
-        energy distance on those sub-samples. This is repeated `chunk_iter`
-        times, and the average is taken. This is a trade-off between speed and
-        accuracy. The larger the chunk size and larger chunk_iter, the more
-        accurate the estimate, but the slower the computation. PTED remains an
-        exact p-value test even when chunking, it simply becomes less sensitive
-        to the difference between x and y. The chunked pted test has time
+        each iteration rather than fully computed. The dataset is divided into
+        sequential chunks of size `chunk_size`; the number of chunks (iterations)
+        is ``max(len(x), len(y)) // chunk_size``, iterating over the larger
+        dataset once and cycling through the smaller one. The average energy
+        distance over all chunks is the final estimate. PTED remains an exact
+        p-value test even when chunking, it simply becomes less sensitive to
+        the difference between x and y. The chunked pted test has time
         complexity O(c^2 * I * D * P), where c is the chunk size, I is the
         number of iterations, D is the number of dimensions, and P is the number
         of permutations. For chunking to be worth it you should have c^2 * I << n^2.
@@ -124,9 +122,10 @@ def pted(
     assert type(x) == type(y), f"x and y must be of the same type, not {type(x)} and {type(y)}"
     assert len(x.shape) >= 2, f"x must be at least 2D, not {x.shape}"
     assert len(y.shape) >= 2, f"y must be at least 2D, not {y.shape}"
-    assert (chunk_size is not None) is (
-        chunk_iter is not None
-    ), "chunk_size and chunk_iter must both be provided for chunked PTED test"
+    if chunk_size is not None:
+        assert (
+            chunk_size <= max(len(x), len(y))
+        ), f"chunk_size ({chunk_size}) must be <= max(len(x), len(y)) ({max(len(x), len(y))})"
     assert (
         x.shape[1:] == y.shape[1:]
     ), f"x and y samples must have the same shape (past first dim), not {x.shape} and {y.shape}"
@@ -142,7 +141,6 @@ def pted(
             permutations=permutations,
             metric=metric,
             chunk_size=int(chunk_size),
-            chunk_iter=int(chunk_iter),
             prog_bar=prog_bar,
         )
     elif is_torch_tensor(x):
@@ -156,7 +154,6 @@ def pted(
             permutations=permutations,
             metric=metric,
             chunk_size=int(chunk_size),
-            chunk_iter=int(chunk_iter),
             prog_bar=prog_bar,
         )
     elif is_jax_array(x):
@@ -168,7 +165,6 @@ def pted(
             permutations=permutations,
             metric=metric,
             chunk_size=int(chunk_size),
-            chunk_iter=int(chunk_iter),
             prog_bar=prog_bar,
         )
     else:
@@ -200,7 +196,6 @@ def pted_coverage_test(
     warn_confidence: Optional[float] = 1e-3,
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     sbc_histogram: Optional[str] = None,
     sbc_bins: Optional[int] = None,
     pit_plot: Optional[str] = None,
@@ -269,10 +264,10 @@ def pted_coverage_test(
             (default: False)
         chunk_size (Optional[int]): If not None, use chunked energy distance
             estimation. This is useful for large datasets. The chunk size is the
-            number of samples to use for each chunk. If None, use the full
-            dataset.
-        chunk_iter (Optional[int]): The chunk iter is the number of iterations
-            to use with the given chunk size.
+            number of samples per chunk. The number of chunks is determined
+            automatically as ``max(len(x), len(y)) // chunk_size``, iterating
+            over the larger dataset once and cycling through the smaller one.
+            If None, use the full dataset.
         sbc_histogram (Optional[str]): If given, the path/filename to save a
             Simulation-Based-Calibration histogram.
         sbc_bins (Optional[int]): If given, force the histogram to have the provided
@@ -295,14 +290,13 @@ def pted_coverage_test(
         samples in x and y, D is the number of dimensions, and P is the number
         of permutations. For large datasets this can get unwieldy, so chunking
         is recommended. For chunking, the energy distance will be estimated at
-        each iteration rather than fully computed. To estimate the energy
-        distance, we take `chunk_size` sub-samples from x and y, and compute the
-        energy distance on those sub-samples. This is repeated `chunk_iter`
-        times, and the average is taken. This is a trade-off between speed and
-        accuracy. The larger the chunk size and larger chunk_iter, the more
-        accurate the estimate, but the slower the computation. PTED remains an
-        exact p-value test even when chunking, it simply becomes less sensitive
-        to the difference between x and y. The chunked pted test has time
+        each iteration rather than fully computed. The dataset is divided into
+        sequential chunks of size `chunk_size`; the number of chunks (iterations)
+        is ``max(len(x), len(y)) // chunk_size``, iterating over the larger
+        dataset once and cycling through the smaller one. The average energy
+        distance over all chunks is the final estimate. PTED remains an exact
+        p-value test even when chunking, it simply becomes less sensitive to
+        the difference between x and y. The chunked pted test has time
         complexity O(c^2 * I * D * P), where c is the chunk size, I is the
         number of iterations, D is the number of dimensions, and P is the number
         of permutations. For chunking to be worth it you should have c^2 * I << n^2.
@@ -328,7 +322,6 @@ def pted_coverage_test(
             return_all=True,
             two_tailed=False,
             chunk_size=chunk_size,
-            chunk_iter=chunk_iter,
         )
         test_stats.append(test)
         permute_stats.append(permute)
diff --git a/src/pted/utils.py b/src/pted/utils.py
index 57df5fb..a239652 100644
--- a/src/pted/utils.py
+++ b/src/pted/utils.py
@@ -90,17 +90,20 @@ def _energy_distance_estimate_numpy(
     x: np.ndarray,
     y: np.ndarray,
     chunk_size: int,
-    chunk_iter: int,
     metric: Union[str, float] = "euclidean",
 ) -> float:
 
+    nx = len(x)
+    ny = len(y)
+    n_iter = max(nx, ny) // chunk_size
+
     E_est = []
-    for _ in range(chunk_iter):
-        # Randomly sample a chunk of data
-        idx = np.random.choice(len(x), size=min(len(x), chunk_size), replace=False)
-        x_chunk = x[idx]
-        idy = np.random.choice(len(y), size=min(len(y), chunk_size), replace=False)
-        y_chunk = y[idy]
+    for i in range(n_iter):
+        # Iterate over the larger dataset, cycling the smaller one
+        x_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % nx
+        y_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % ny
+        x_chunk = x[x_idx]
+        y_chunk = y[y_idx]
 
         # Compute the energy distance
         E_est.append(_energy_distance_numpy(x_chunk, y_chunk, metric=metric))
@@ -111,17 +114,20 @@ def _energy_distance_estimate_torch(
     x: torch.Tensor,
     y: torch.Tensor,
     chunk_size: int,
-    chunk_iter: int,
     metric: Union[str, float] = "euclidean",
 ) -> float:
 
+    nx = len(x)
+    ny = len(y)
+    n_iter = max(nx, ny) // chunk_size
+
     E_est = []
-    for _ in range(chunk_iter):
-        # Randomly sample a chunk of data
-        idx = np.random.choice(len(x), size=min(len(x), chunk_size), replace=False)
-        x_chunk = x[torch.tensor(idx)]
-        idy = np.random.choice(len(y), size=min(len(y), chunk_size), replace=False)
-        y_chunk = y[torch.tensor(idy)]
+    for i in range(n_iter):
+        # Iterate over the larger dataset, cycling the smaller one
+        x_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % nx
+        y_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % ny
+        x_chunk = x[torch.tensor(x_idx)]
+        y_chunk = y[torch.tensor(y_idx)]
 
         # Compute the energy distance
         E_est.append(_energy_distance_torch(x_chunk, y_chunk, metric=metric))
@@ -154,17 +160,20 @@ def _energy_distance_estimate_jax(
     x,
     y,
     chunk_size: int,
-    chunk_iter: int,
     metric: Union[str, float] = "euclidean",
 ) -> float:
 
+    nx = len(x)
+    ny = len(y)
+    n_iter = max(nx, ny) // chunk_size
+
     E_est = []
-    for _ in range(chunk_iter):
-        # Randomly sample a chunk of data
-        idx = np.random.choice(len(x), size=min(len(x), chunk_size), replace=False)
-        x_chunk = x[idx]
-        idy = np.random.choice(len(y), size=min(len(y), chunk_size), replace=False)
-        y_chunk = y[idy]
+    for i in range(n_iter):
+        # Iterate over the larger dataset, cycling the smaller one
+        x_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % nx
+        y_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % ny
+        x_chunk = x[x_idx]
+        y_chunk = y[y_idx]
 
         # Compute the energy distance
         E_est.append(_energy_distance_jax(x_chunk, y_chunk, metric=metric))
@@ -177,20 +186,19 @@ def pted_chunk_numpy(
     permutations: int = 100,
     metric: str = "euclidean",
     chunk_size: int = 100,
-    chunk_iter: int = 10,
     prog_bar: bool = False,
 ) -> tuple[float, list[float]]:
     assert np.all(np.isfinite(x)) and np.all(np.isfinite(y)), "Input contains NaN or Inf!"
     nx = len(x)
 
-    test_stat = _energy_distance_estimate_numpy(x, y, chunk_size, chunk_iter, metric=metric)
+    test_stat = _energy_distance_estimate_numpy(x, y, chunk_size, metric=metric)
     permute_stats = []
     for _ in trange(permutations, disable=not prog_bar):
         z = np.concatenate((x, y), axis=0)
         z = z[np.random.permutation(len(z))]
         x, y = z[:nx], z[nx:]
         permute_stats.append(
-            _energy_distance_estimate_numpy(x, y, chunk_size, chunk_iter, metric=metric)
+            _energy_distance_estimate_numpy(x, y, chunk_size, metric=metric)
         )
     return test_stat, permute_stats
 
@@ -201,7 +209,6 @@ def pted_chunk_torch(
     permutations: int = 100,
     metric: Union[str, float] = "euclidean",
     chunk_size: int = 100,
-    chunk_iter: int = 10,
     prog_bar: bool = False,
 ) -> tuple[float, list[float]]:
     assert torch.__version__ != "null", "PyTorch is not installed! try: `pip install torch`"
@@ -210,14 +217,14 @@ def pted_chunk_torch(
     ), "Input contains NaN or Inf!"
     nx = len(x)
 
-    test_stat = _energy_distance_estimate_torch(x, y, chunk_size, chunk_iter, metric=metric)
+    test_stat = _energy_distance_estimate_torch(x, y, chunk_size, metric=metric)
     permute_stats = []
     for _ in trange(permutations, disable=not prog_bar):
         z = torch.cat((x, y), dim=0)
         z = z[torch.randperm(len(z))]
         x, y = z[:nx], z[nx:]
         permute_stats.append(
-            _energy_distance_estimate_torch(x, y, chunk_size, chunk_iter, metric=metric)
+            _energy_distance_estimate_torch(x, y, chunk_size, metric=metric)
         )
     return test_stat, permute_stats
 
@@ -309,21 +316,20 @@ def pted_chunk_jax(
     permutations: int = 100,
     metric: Union[str, float] = "euclidean",
     chunk_size: int = 100,
-    chunk_iter: int = 10,
     prog_bar: bool = False,
 ) -> tuple[float, list[float]]:
     assert jax is not None, "JAX is not installed! try: `pip install jax`"
     assert jnp.all(jnp.isfinite(x)) and jnp.all(jnp.isfinite(y)), "Input contains NaN or Inf!"
     nx = len(x)
 
-    test_stat = _energy_distance_estimate_jax(x, y, chunk_size, chunk_iter, metric=metric)
+    test_stat = _energy_distance_estimate_jax(x, y, chunk_size, metric=metric)
     permute_stats = []
     for _ in trange(permutations, disable=not prog_bar):
         z = jnp.concatenate([x, y], axis=0)
         z = z[np.random.permutation(len(z))]
         x, y = z[:nx], z[nx:]
         permute_stats.append(
-            _energy_distance_estimate_jax(x, y, chunk_size, chunk_iter, metric=metric)
+            _energy_distance_estimate_jax(x, y, chunk_size, metric=metric)
         )
     return test_stat, permute_stats
 
diff --git a/tests/test_pted.py b/tests/test_pted.py
index 35f397a..12c8deb 100644
--- a/tests/test_pted.py
+++ b/tests/test_pted.py
@@ -104,11 +104,11 @@ def test_pted_chunk_torch():
     D = 10
     x = torch.randn(1000, D)
     y = torch.randn(1000, D)
-    p = pted.pted(x, y, chunk_size=100, chunk_iter=10)
+    p = pted.pted(x, y, chunk_size=100)
     assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))"
 
     y = torch.rand(1000, D)
-    p = pted.pted(x, y, chunk_size=100, chunk_iter=10)
+    p = pted.pted(x, y, chunk_size=100)
     assert p < 1e-2, f"p-value {p} is not in the expected range (~0)"
 
 
@@ -119,11 +119,27 @@ def test_pted_chunk_numpy():
     D = 10
     x = np.random.normal(size=(1000, D))
     y = np.random.normal(size=(1000, D))
-    p = pted.pted(x, y, chunk_size=100, chunk_iter=10)
+    p = pted.pted(x, y, chunk_size=100)
     assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))"
 
     y = np.random.uniform(size=(1000, D))
-    p = pted.pted(x, y, chunk_size=100, chunk_iter=10)
+    p = pted.pted(x, y, chunk_size=100)
+    assert p < 1e-2, f"p-value {p} is not in the expected range (~0)"
+
+
+def test_pted_chunk_mismatched_sizes():
+    """Chunked PTED correctly cycles the smaller dataset for mismatched x/y sizes."""
+    np.random.seed(0)
+    D = 5
+    # x has 800 samples, y has 300 samples; chunk_size=100 → 8 iterations, cycling y
+    x = np.random.normal(size=(800, D))
+    y = np.random.normal(size=(300, D))
+    p = pted.pted(x, y, chunk_size=100)
+    assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))"
+
+    # Different distributions should give small p-value even with mismatched sizes
+    y_diff = np.random.uniform(size=(300, D))
+    p = pted.pted(x, y_diff, chunk_size=100)
     assert p < 1e-2, f"p-value {p} is not in the expected range (~0)"
 
 
@@ -225,17 +241,17 @@ def test_pted_jax():
 def test_pted_chunk_jax():
     if jax is None:
         pytest.skip("jax not installed")
-    np.random.seed(42)
+    np.random.seed(0)
 
     # example 2 sample test
     D = 3
-    x = jnp.array(np.random.normal(size=(100, D)))
-    y = jnp.array(np.random.normal(size=(100, D)))
-    p = pted.pted(x, y, chunk_size=100, chunk_iter=10)
+    x = jnp.array(np.random.normal(size=(200, D)))
+    y = jnp.array(np.random.normal(size=(200, D)))
+    p = pted.pted(x, y, chunk_size=50)
     assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))"
 
-    y = jnp.array(np.random.uniform(size=(110, D)))
-    p = pted.pted(x, y, chunk_size=100, chunk_iter=10)
+    y = jnp.array(np.random.uniform(size=(300, D)))
+    p = pted.pted(x, y, chunk_size=50)
     assert p < 1e-2, f"p-value {p} is not in the expected range (~0)"
 
 
@@ -311,7 +327,7 @@ def test_energy_distance_estimate_jax():
     np.random.seed(0)
     x = jnp.array(np.random.normal(size=(100, 4)))
     y = jnp.array(np.random.normal(size=(100, 4)))
-    ed = pted.utils._energy_distance_estimate_jax(x, y, chunk_size=20, chunk_iter=5)
+    ed = pted.utils._energy_distance_estimate_jax(x, y, chunk_size=20)
     assert np.isfinite(ed)
 
 
@@ -391,7 +407,7 @@ def test_energy_distance_numpy_torch_jax_agree():
 
 
 def test_energy_distance_estimate_numpy_torch_jax_agree():
-    """_energy_distance_estimate_{numpy,torch,jax} return close values for the same seed/data."""
+    """_energy_distance_estimate_{numpy,torch,jax} return close values for the same data."""
     if torch is None:
         pytest.skip("torch not installed")
     if jax is None:
@@ -402,16 +418,13 @@ def test_energy_distance_estimate_numpy_torch_jax_agree():
     x_np = np.random.normal(size=(200, 5)).astype(np.float32)
     y_np = np.random.normal(size=(200, 5)).astype(np.float32)
 
-    # Run with the same seed so the same chunks are sampled
-    np.random.seed(0)
-    ed_numpy = pted.utils._energy_distance_estimate_numpy(x_np, y_np, chunk_size=50, chunk_iter=5)
-    np.random.seed(0)
+    # Deterministic sequential chunking: no seed needed
+    ed_numpy = pted.utils._energy_distance_estimate_numpy(x_np, y_np, chunk_size=50)
     ed_torch = pted.utils._energy_distance_estimate_torch(
-        torch.tensor(x_np), torch.tensor(y_np), chunk_size=50, chunk_iter=5
+        torch.tensor(x_np), torch.tensor(y_np), chunk_size=50
     )
-    np.random.seed(0)
     ed_jax = pted.utils._energy_distance_estimate_jax(
-        jnp.array(x_np), jnp.array(y_np), chunk_size=50, chunk_iter=5
+        jnp.array(x_np), jnp.array(y_np), chunk_size=50
     )
 
     assert ed_numpy == pytest.approx(

From 67d6552fc0f26feea0510421177fd9a8e3ef1317 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 1 Aug 2026 10:48:13 +0000
Subject: [PATCH 03/12] refactor: consolidate energy distance estimate
 functions, use slicing, add chunk_size fallback

- Replace three `_energy_distance_estimate_{numpy,torch,jax}` functions with a
  single `_energy_distance_estimate` core that uses plain slicing via a
  `_tile_to_length` helper; the named functions become thin wrappers for API
  compatibility.
- Removes the `torch.tensor(idx)` fancy-indexing workaround from the torch path.
- Add fallback in `pted()`: when `chunk_size >= len(x)` and `chunk_size >= len(y)`,
  chunking covers the full dataset so regular PTED is used automatically.
- Add `test_pted_chunk_size_fallback` to cover the new fallback behaviour.
---
 src/pted/pted.py   |  3 ++
 src/pted/utils.py  | 97 +++++++++++++++++++++++-----------------------
 tests/test_pted.py | 17 ++++++++
 3 files changed, 69 insertions(+), 48 deletions(-)

diff --git a/src/pted/pted.py b/src/pted/pted.py
index f30aa58..9fd120f 100644
--- a/src/pted/pted.py
+++ b/src/pted/pted.py
@@ -126,6 +126,9 @@ def pted(
         assert (
             chunk_size <= max(len(x), len(y))
         ), f"chunk_size ({chunk_size}) must be <= max(len(x), len(y)) ({max(len(x), len(y))})"
+        # If chunk_size covers both full datasets, chunking adds no benefit
+        if chunk_size >= len(x) and chunk_size >= len(y):
+            chunk_size = None
     assert (
         x.shape[1:] == y.shape[1:]
     ), f"x and y samples must have the same shape (past first dim), not {x.shape} and {y.shape}"
diff --git a/src/pted/utils.py b/src/pted/utils.py
index a239652..ff6a9bf 100644
--- a/src/pted/utils.py
+++ b/src/pted/utils.py
@@ -1,3 +1,4 @@
+import math
 from typing import Union
 from warnings import warn
 
@@ -86,52 +87,67 @@ def _energy_distance_torch(
     return _energy_distance_precompute(D, nx, ny).item()
 
 
-def _energy_distance_estimate_numpy(
-    x: np.ndarray,
-    y: np.ndarray,
+def _tile_to_length(arr, n: int):
+    """Return arr tiled along axis 0 to exactly n rows.
+
+    If len(arr) >= n the array is simply sliced; otherwise it is repeated
+    along axis 0 enough times and then sliced to exactly n rows.
+    """
+    if len(arr) >= n:
+        return arr[:n]
+    reps = math.ceil(n / len(arr))
+    reps_tuple = (reps,) + (1,) * (arr.ndim - 1)
+    if is_torch_tensor(arr):
+        return arr.repeat(reps_tuple)[:n]
+    elif is_jax_array(arr):
+        return jnp.tile(arr, reps_tuple)[:n]
+    else:
+        return np.tile(arr, reps_tuple)[:n]
+
+
+def _energy_distance_estimate(
+    x,
+    y,
     chunk_size: int,
-    metric: Union[str, float] = "euclidean",
+    metric: Union[str, float],
+    energy_distance_fn,
 ) -> float:
+    """Estimate energy distance by averaging over sequential sliced chunks.
 
-    nx = len(x)
-    ny = len(y)
-    n_iter = max(nx, ny) // chunk_size
-
+    Iterates ``max(len(x), len(y)) // chunk_size`` times, using plain slicing
+    on both arrays.  The smaller of the two is tiled along axis 0 as needed so
+    that both arrays are at least ``n_iter * chunk_size`` rows long before the
+    loop begins.
+    """
+    n_iter = max(len(x), len(y)) // chunk_size
+    n_total = n_iter * chunk_size
+    x = _tile_to_length(x, n_total)
+    y = _tile_to_length(y, n_total)
     E_est = []
     for i in range(n_iter):
-        # Iterate over the larger dataset, cycling the smaller one
-        x_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % nx
-        y_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % ny
-        x_chunk = x[x_idx]
-        y_chunk = y[y_idx]
-
-        # Compute the energy distance
-        E_est.append(_energy_distance_numpy(x_chunk, y_chunk, metric=metric))
+        start = i * chunk_size
+        x_chunk = x[start : start + chunk_size]
+        y_chunk = y[start : start + chunk_size]
+        E_est.append(energy_distance_fn(x_chunk, y_chunk, metric=metric))
     return np.mean(E_est)
 
 
+def _energy_distance_estimate_numpy(
+    x: np.ndarray,
+    y: np.ndarray,
+    chunk_size: int,
+    metric: Union[str, float] = "euclidean",
+) -> float:
+    return _energy_distance_estimate(x, y, chunk_size, metric, _energy_distance_numpy)
+
+
 def _energy_distance_estimate_torch(
     x: torch.Tensor,
     y: torch.Tensor,
     chunk_size: int,
     metric: Union[str, float] = "euclidean",
 ) -> float:
-
-    nx = len(x)
-    ny = len(y)
-    n_iter = max(nx, ny) // chunk_size
-
-    E_est = []
-    for i in range(n_iter):
-        # Iterate over the larger dataset, cycling the smaller one
-        x_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % nx
-        y_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % ny
-        x_chunk = x[torch.tensor(x_idx)]
-        y_chunk = y[torch.tensor(y_idx)]
-
-        # Compute the energy distance
-        E_est.append(_energy_distance_torch(x_chunk, y_chunk, metric=metric))
-    return np.mean(E_est)
+    return _energy_distance_estimate(x, y, chunk_size, metric, _energy_distance_torch)
 
 
 def _jax_cdist(x, y, p: float = 2.0):
@@ -162,22 +178,7 @@ def _energy_distance_estimate_jax(
     chunk_size: int,
     metric: Union[str, float] = "euclidean",
 ) -> float:
-
-    nx = len(x)
-    ny = len(y)
-    n_iter = max(nx, ny) // chunk_size
-
-    E_est = []
-    for i in range(n_iter):
-        # Iterate over the larger dataset, cycling the smaller one
-        x_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % nx
-        y_idx = np.arange(i * chunk_size, (i + 1) * chunk_size) % ny
-        x_chunk = x[x_idx]
-        y_chunk = y[y_idx]
-
-        # Compute the energy distance
-        E_est.append(_energy_distance_jax(x_chunk, y_chunk, metric=metric))
-    return np.mean(E_est)
+    return _energy_distance_estimate(x, y, chunk_size, metric, _energy_distance_jax)
 
 
 def pted_chunk_numpy(
diff --git a/tests/test_pted.py b/tests/test_pted.py
index 12c8deb..aa55900 100644
--- a/tests/test_pted.py
+++ b/tests/test_pted.py
@@ -143,6 +143,23 @@ def test_pted_chunk_mismatched_sizes():
     assert p < 1e-2, f"p-value {p} is not in the expected range (~0)"
 
 
+def test_pted_chunk_size_fallback():
+    """chunk_size >= both dataset sizes falls back to regular (non-chunked) PTED."""
+    np.random.seed(7)
+    D = 5
+    x = np.random.normal(size=(50, D))
+    y = np.random.normal(size=(50, D))
+
+    # chunk_size equals both lengths → should silently fall back to regular PTED
+    p_chunk = pted.pted(x, y, chunk_size=50)
+    # Rerun with explicit non-chunked PTED using same seed for reference
+    np.random.seed(7)
+    p_plain = pted.pted(x, y)
+    # Both should give a valid p-value in U(0,1)
+    assert p_chunk > 1e-2 and p_chunk < 0.99, f"p-value {p_chunk} outside expected range"
+    assert p_plain > 1e-2 and p_plain < 0.99, f"p-value {p_plain} outside expected range"
+
+
 def test_pted_coverage_edgecase():
     # Test with single simulation
     g = np.random.normal(size=(1, 10))

From 12f4d884b8b556bfecae9db46a4e48f5eb5b9634 Mon Sep 17 00:00:00 2001
From: Connor Stone <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 08:03:50 -0400
Subject: [PATCH 04/12] refactor for cleaner energy distance estimate
 implementation

---
 src/pted/pted.py   |  3 --
 src/pted/utils.py  | 96 +++++++++++++++++++---------------------------
 tests/test_pted.py | 60 ++++++++---------------------
 3 files changed, 54 insertions(+), 105 deletions(-)

diff --git a/src/pted/pted.py b/src/pted/pted.py
index 9fd120f..f1e36f9 100644
--- a/src/pted/pted.py
+++ b/src/pted/pted.py
@@ -123,9 +123,6 @@ def pted(
     assert len(x.shape) >= 2, f"x must be at least 2D, not {x.shape}"
     assert len(y.shape) >= 2, f"y must be at least 2D, not {y.shape}"
     if chunk_size is not None:
-        assert (
-            chunk_size <= max(len(x), len(y))
-        ), f"chunk_size ({chunk_size}) must be <= max(len(x), len(y)) ({max(len(x), len(y))})"
         # If chunk_size covers both full datasets, chunking adds no benefit
         if chunk_size >= len(x) and chunk_size >= len(y):
             chunk_size = None
diff --git a/src/pted/utils.py b/src/pted/utils.py
index ff6a9bf..6d99d28 100644
--- a/src/pted/utils.py
+++ b/src/pted/utils.py
@@ -1,4 +1,3 @@
-import math
 from typing import Union
 from warnings import warn
 
@@ -87,22 +86,27 @@ def _energy_distance_torch(
     return _energy_distance_precompute(D, nx, ny).item()
 
 
-def _tile_to_length(arr, n: int):
-    """Return arr tiled along axis 0 to exactly n rows.
+def _chunk_slices(lenx: int, leny: int, chunk_size: int):
+    """Yield slices for chunking two arrays of lengths lenx and leny.
 
-    If len(arr) >= n the array is simply sliced; otherwise it is repeated
-    along axis 0 enough times and then sliced to exactly n rows.
+    The smaller of the two is tiled along axis 0 as needed so that both arrays
+    are at least ``n_iter * chunk_size`` rows long before the loop begins.
     """
-    if len(arr) >= n:
-        return arr[:n]
-    reps = math.ceil(n / len(arr))
-    reps_tuple = (reps,) + (1,) * (arr.ndim - 1)
-    if is_torch_tensor(arr):
-        return arr.repeat(reps_tuple)[:n]
-    elif is_jax_array(arr):
-        return jnp.tile(arr, reps_tuple)[:n]
-    else:
-        return np.tile(arr, reps_tuple)[:n]
+    nx = max(1, lenx // chunk_size)
+    x_slices = []
+    for i in range(nx):
+        start = i * chunk_size
+        x_slices.append(slice(start, start + chunk_size))
+
+    y_slices = []
+    ny = max(1, leny // chunk_size)
+    for i in range(ny):
+        start = i * chunk_size
+        y_slices.append(slice(start, start + chunk_size))
+
+    n_iter = max(lenx, leny) // chunk_size
+    for i in range(n_iter):
+        yield x_slices[i % nx], y_slices[i % ny]
 
 
 def _energy_distance_estimate(
@@ -119,37 +123,12 @@ def _energy_distance_estimate(
     that both arrays are at least ``n_iter * chunk_size`` rows long before the
     loop begins.
     """
-    n_iter = max(len(x), len(y)) // chunk_size
-    n_total = n_iter * chunk_size
-    x = _tile_to_length(x, n_total)
-    y = _tile_to_length(y, n_total)
     E_est = []
-    for i in range(n_iter):
-        start = i * chunk_size
-        x_chunk = x[start : start + chunk_size]
-        y_chunk = y[start : start + chunk_size]
-        E_est.append(energy_distance_fn(x_chunk, y_chunk, metric=metric))
+    for cx, cy in _chunk_slices(len(x), len(y), chunk_size):
+        E_est.append(energy_distance_fn(x[cx], y[cy], metric=metric))
     return np.mean(E_est)
 
 
-def _energy_distance_estimate_numpy(
-    x: np.ndarray,
-    y: np.ndarray,
-    chunk_size: int,
-    metric: Union[str, float] = "euclidean",
-) -> float:
-    return _energy_distance_estimate(x, y, chunk_size, metric, _energy_distance_numpy)
-
-
-def _energy_distance_estimate_torch(
-    x: torch.Tensor,
-    y: torch.Tensor,
-    chunk_size: int,
-    metric: Union[str, float] = "euclidean",
-) -> float:
-    return _energy_distance_estimate(x, y, chunk_size, metric, _energy_distance_torch)
-
-
 def _jax_cdist(x, y, p: float = 2.0):
     if p == 2.0:
         # Squared-norm identity avoids materializing the (nx, ny, d) diff tensor.
@@ -172,15 +151,6 @@ def _energy_distance_jax(x, y, metric: Union[str, float] = "euclidean") -> float
     return float(_energy_distance_precompute(D, nx, ny))
 
 
-def _energy_distance_estimate_jax(
-    x,
-    y,
-    chunk_size: int,
-    metric: Union[str, float] = "euclidean",
-) -> float:
-    return _energy_distance_estimate(x, y, chunk_size, metric, _energy_distance_jax)
-
-
 def pted_chunk_numpy(
     x: np.ndarray,
     y: np.ndarray,
@@ -192,14 +162,18 @@ def pted_chunk_numpy(
     assert np.all(np.isfinite(x)) and np.all(np.isfinite(y)), "Input contains NaN or Inf!"
     nx = len(x)
 
-    test_stat = _energy_distance_estimate_numpy(x, y, chunk_size, metric=metric)
+    test_stat = _energy_distance_estimate(
+        x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_numpy
+    )
     permute_stats = []
     for _ in trange(permutations, disable=not prog_bar):
         z = np.concatenate((x, y), axis=0)
         z = z[np.random.permutation(len(z))]
         x, y = z[:nx], z[nx:]
         permute_stats.append(
-            _energy_distance_estimate_numpy(x, y, chunk_size, metric=metric)
+            _energy_distance_estimate(
+                x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_numpy
+            )
         )
     return test_stat, permute_stats
 
@@ -218,14 +192,18 @@ def pted_chunk_torch(
     ), "Input contains NaN or Inf!"
     nx = len(x)
 
-    test_stat = _energy_distance_estimate_torch(x, y, chunk_size, metric=metric)
+    test_stat = _energy_distance_estimate(
+        x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_torch
+    )
     permute_stats = []
     for _ in trange(permutations, disable=not prog_bar):
         z = torch.cat((x, y), dim=0)
         z = z[torch.randperm(len(z))]
         x, y = z[:nx], z[nx:]
         permute_stats.append(
-            _energy_distance_estimate_torch(x, y, chunk_size, metric=metric)
+            _energy_distance_estimate(
+                x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_torch
+            )
         )
     return test_stat, permute_stats
 
@@ -323,14 +301,18 @@ def pted_chunk_jax(
     assert jnp.all(jnp.isfinite(x)) and jnp.all(jnp.isfinite(y)), "Input contains NaN or Inf!"
     nx = len(x)
 
-    test_stat = _energy_distance_estimate_jax(x, y, chunk_size, metric=metric)
+    test_stat = _energy_distance_estimate(
+        x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_jax
+    )
     permute_stats = []
     for _ in trange(permutations, disable=not prog_bar):
         z = jnp.concatenate([x, y], axis=0)
         z = z[np.random.permutation(len(z))]
         x, y = z[:nx], z[nx:]
         permute_stats.append(
-            _energy_distance_estimate_jax(x, y, chunk_size, metric=metric)
+            _energy_distance_estimate(
+                x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_jax
+            )
         )
     return test_stat, permute_stats
 
diff --git a/tests/test_pted.py b/tests/test_pted.py
index aa55900..6e4605e 100644
--- a/tests/test_pted.py
+++ b/tests/test_pted.py
@@ -282,11 +282,6 @@ def test_pted_coverage_jax():
     assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))"
 
 
-# ---------------------------------------------------------------------------
-# Unit tests for newly-added utils functions
-# ---------------------------------------------------------------------------
-
-
 def test_is_jax_array_with_jax():
     """is_jax_array returns True for a real JAX array and False for other types."""
     if jax is None:
@@ -337,17 +332,6 @@ def test_energy_distance_jax():
     assert abs(ed) < 1e-6
 
 
-def test_energy_distance_estimate_jax():
-    """_energy_distance_estimate_jax returns a finite scalar."""
-    if jax is None:
-        pytest.skip("jax not installed")
-    np.random.seed(0)
-    x = jnp.array(np.random.normal(size=(100, 4)))
-    y = jnp.array(np.random.normal(size=(100, 4)))
-    ed = pted.utils._energy_distance_estimate_jax(x, y, chunk_size=20)
-    assert np.isfinite(ed)
-
-
 def test_pted_jax_no_jax(monkeypatch):
     """pted_jax raises AssertionError when JAX is not installed."""
     monkeypatch.setattr("pted.utils.jax", None)
@@ -378,11 +362,6 @@ def test_pted_chunk_torch_no_torch(monkeypatch):
         pted.utils.pted_chunk_torch(np.zeros((5, 2)), np.zeros((5, 2)))
 
 
-# ---------------------------------------------------------------------------
-# Cross-backend consistency tests
-# ---------------------------------------------------------------------------
-
-
 def test_jax_cdist_matches_scipy():
     """_jax_cdist (L2) and scipy cdist produce the same pairwise distances."""
     if jax is None:
@@ -423,30 +402,21 @@ def test_energy_distance_numpy_torch_jax_agree():
     ), f"numpy ({ed_numpy}) and jax ({ed_jax}) energy distances differ"
 
 
-def test_energy_distance_estimate_numpy_torch_jax_agree():
-    """_energy_distance_estimate_{numpy,torch,jax} return close values for the same data."""
-    if torch is None:
-        pytest.skip("torch not installed")
-    if jax is None:
-        pytest.skip("jax not installed")
-
+def test_energy_distance_estimate_matches_energy_distance():
+    """_energy_distance_estimate returns the same value as _energy_distance for small inputs."""
     np.random.seed(123)
-    # Use float32 so all backends operate at the same precision
-    x_np = np.random.normal(size=(200, 5)).astype(np.float32)
-    y_np = np.random.normal(size=(200, 5)).astype(np.float32)
-
-    # Deterministic sequential chunking: no seed needed
-    ed_numpy = pted.utils._energy_distance_estimate_numpy(x_np, y_np, chunk_size=50)
-    ed_torch = pted.utils._energy_distance_estimate_torch(
-        torch.tensor(x_np), torch.tensor(y_np), chunk_size=50
-    )
-    ed_jax = pted.utils._energy_distance_estimate_jax(
-        jnp.array(x_np), jnp.array(y_np), chunk_size=50
+    x = np.random.normal(size=(100, 5))
+    y = np.random.uniform(size=(100, 5))
+
+    ed_direct = pted.utils._energy_distance_numpy(x, y)
+    ed_estimate = pted.utils._energy_distance_estimate(
+        x,
+        y,
+        chunk_size=25,
+        metric="euclidean",
+        energy_distance_fn=pted.utils._energy_distance_numpy,
     )
 
-    assert ed_numpy == pytest.approx(
-        ed_torch, rel=1e-4
-    ), f"numpy ({ed_numpy}) and torch ({ed_torch}) energy distance estimates differ"
-    assert ed_numpy == pytest.approx(
-        ed_jax, rel=1e-4
-    ), f"numpy ({ed_numpy}) and jax ({ed_jax}) energy distance estimates differ"
+    assert ed_direct == pytest.approx(
+        ed_estimate, rel=1.5e-1
+    ), f"direct ({ed_direct}) and estimate ({ed_estimate}) energy distances differ"

From 0b28de29c1b86bf6b39e0097ebb098f040e0590f Mon Sep 17 00:00:00 2001
From: Connor Stone <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 08:23:23 -0400
Subject: [PATCH 05/12] minor extra comments in README

---
 README.md | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 9b6c033..4b16eb1 100644
--- a/README.md
+++ b/README.md
@@ -365,10 +365,12 @@ iterating over the larger dataset once and cycling through the smaller one if
 their sizes differ. The larger the chunk size, the closer the estimate will be
 to the true energy distance, but it will take more compute.
 
-Note that the computational complexity for standard PTED goes as 
-`O((n_samp_x + n_samp_y)^2)` while the chunked version goes as 
-`O(n_iter * (2 * chunk_size)^2)` where `n_iter = max(n_samp_x, n_samp_y) // chunk_size`,
-so plan your chunking accordingly.
+Note that the computational complexity for standard PTED goes as `O((n_samp_x +
+n_samp_y)^2)` while the chunked version goes as `O(n_iter * (2 * chunk_size)^2)`
+where `n_iter = max(n_samp_x, n_samp_y) // chunk_size`, so plan your chunking
+accordingly. For a given chunk size, the computational complexity of PTED grows
+linearly with dataset size, much like other large scale (machine learning
+oriented) two sample tests.
 
 Example:
 ```python

From fb3a46897c4b250f0246de5d0e622e390fe29776 Mon Sep 17 00:00:00 2001
From: "Connor Stone, PhD" <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 08:41:10 -0400
Subject: [PATCH 06/12] assert chunk_size is a strictly positive integer

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 src/pted/pted.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/pted/pted.py b/src/pted/pted.py
index f1e36f9..590941c 100644
--- a/src/pted/pted.py
+++ b/src/pted/pted.py
@@ -123,6 +123,7 @@ def pted(
     assert len(x.shape) >= 2, f"x must be at least 2D, not {x.shape}"
     assert len(y.shape) >= 2, f"y must be at least 2D, not {y.shape}"
     if chunk_size is not None:
+        assert chunk_size > 0, "chunk_size must be > 0"
         # If chunk_size covers both full datasets, chunking adds no benefit
         if chunk_size >= len(x) and chunk_size >= len(y):
             chunk_size = None

From c0f096816b7b10910e5fc491141577fd923e1cdf Mon Sep 17 00:00:00 2001
From: "Connor Stone, PhD" <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 08:42:20 -0400
Subject: [PATCH 07/12] fix wording on test to be more accurate

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 tests/test_pted.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_pted.py b/tests/test_pted.py
index 6e4605e..bb3c925 100644
--- a/tests/test_pted.py
+++ b/tests/test_pted.py
@@ -403,7 +403,7 @@ def test_energy_distance_numpy_torch_jax_agree():
 
 
 def test_energy_distance_estimate_matches_energy_distance():
-    """_energy_distance_estimate returns the same value as _energy_distance for small inputs."""
+    """_energy_distance_estimate returns a value close to _energy_distance for small inputs."""
     np.random.seed(123)
     x = np.random.normal(size=(100, 5))
     y = np.random.uniform(size=(100, 5))

From 4b7cb2b1542159fa853abdd313cd33006c13de94 Mon Sep 17 00:00:00 2001
From: "Connor Stone, PhD" <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 08:43:28 -0400
Subject: [PATCH 08/12] explain fallback behaviour for chunk_size larger than
 dataset size

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 4b16eb1..56a63db 100644
--- a/README.md
+++ b/README.md
@@ -285,7 +285,7 @@ def pted(
 * **permutations** *(int)*: number of permutations to run. This determines how accurately the p-value is computed.
 * **metric** *(Union[str, float])*: distance metric to use. See scipy.spatial.distance.cdist for the list of available metrics with numpy. See torch.cdist when using PyTorch, note that the metric is passed as the "p" for torch.cdist and therefore is a float from 0 to inf. When using JAX arrays, the metric is passed as the "ord" for jnp.linalg.norm and therefore is also a float.
 * **return_all** *(bool)*: if True, return the test statistic and the permuted statistics with the p-value. If False, just return the p-value. bool (default: False)
-* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. The chunk size is the number of samples per chunk. The number of chunks is determined automatically as `max(len(x), len(y)) // chunk_size`, iterating over the larger dataset once and cycling through the smaller one. If None, use the full dataset.
+* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. The chunk size is the number of samples per chunk. The number of chunks is determined automatically as `max(len(x), len(y)) // chunk_size`, iterating over the larger dataset once and cycling through the smaller one. If `chunk_size >= len(x)` and `chunk_size >= len(y)`, PTED falls back to the full (non-chunked) computation. If None, use the full dataset.
 * **two_tailed** *(bool)*: if True, compute a two-tailed p-value. This is useful if you want to reject the null hypothesis when x and y are either too similar or too different. If False, only checks for dissimilarity but is more sensitive. Default is True.
 * **prog_bar** *(bool)*: if True, show a progress bar to track the progress of permutation tests. Default is False.
 

From 27ef2dd0b29b44cfa604676ead54b0132afab065 Mon Sep 17 00:00:00 2001
From: Connor Stone <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 09:24:19 -0400
Subject: [PATCH 09/12] refactor _chunk_slices util for better performance

---
 src/pted/utils.py | 22 ++++++++--------------
 1 file changed, 8 insertions(+), 14 deletions(-)

diff --git a/src/pted/utils.py b/src/pted/utils.py
index 6d99d28..e441ad3 100644
--- a/src/pted/utils.py
+++ b/src/pted/utils.py
@@ -89,24 +89,18 @@ def _energy_distance_torch(
 def _chunk_slices(lenx: int, leny: int, chunk_size: int):
     """Yield slices for chunking two arrays of lengths lenx and leny.
 
-    The smaller of the two is tiled along axis 0 as needed so that both arrays
-    are at least ``n_iter * chunk_size`` rows long before the loop begins.
+    The smaller of the two is cycled while the larger is iterated through
+    completely (minus the last incomplete chunk if any).
     """
     nx = max(1, lenx // chunk_size)
-    x_slices = []
-    for i in range(nx):
-        start = i * chunk_size
-        x_slices.append(slice(start, start + chunk_size))
-
-    y_slices = []
     ny = max(1, leny // chunk_size)
-    for i in range(ny):
-        start = i * chunk_size
-        y_slices.append(slice(start, start + chunk_size))
 
-    n_iter = max(lenx, leny) // chunk_size
-    for i in range(n_iter):
-        yield x_slices[i % nx], y_slices[i % ny]
+    for i in range(max(nx, ny)):
+        ix = i % nx
+        iy = i % ny
+        yield slice(ix * chunk_size, (ix + 1) * chunk_size), slice(
+            iy * chunk_size, (iy + 1) * chunk_size
+        )
 
 
 def _energy_distance_estimate(

From 49d4fcc0c71a1578842ead07f491f6e3a79d6896 Mon Sep 17 00:00:00 2001
From: Connor Stone <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 10:54:14 -0400
Subject: [PATCH 10/12] more stable tests, slight efficiency improvement for
 chunking pted

---
 src/pted/utils.py  | 25 +++++++++----------------
 tests/test_pted.py | 10 ++++++++++
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/pted/utils.py b/src/pted/utils.py
index e441ad3..17ce375 100644
--- a/src/pted/utils.py
+++ b/src/pted/utils.py
@@ -19,9 +19,11 @@ class torch:
 try:
     import jax
     import jax.numpy as jnp
+    from jax import jit
 except ImportError:
     jax = None
     jnp = None
+    jit = lambda *a, **k: lambda f: f  # type: ignore
 
 
 __all__ = (
@@ -123,14 +125,8 @@ def _energy_distance_estimate(
     return np.mean(E_est)
 
 
+@jit(static_argnames=["p"])
 def _jax_cdist(x, y, p: float = 2.0):
-    if p == 2.0:
-        # Squared-norm identity avoids materializing the (nx, ny, d) diff tensor.
-        # ||x_i - y_j||^2 = ||x_i||^2 + ||y_j||^2 - 2 * x_i . y_j
-        x_sq = jnp.sum(x**2, axis=-1)  # (nx,)
-        y_sq = jnp.sum(y**2, axis=-1)  # (ny,)
-        sq_dist = x_sq[:, None] + y_sq[None, :] - 2.0 * (x @ y.T)
-        return jnp.sqrt(jnp.maximum(sq_dist, 0.0))
     # For general p-norms use vmap to avoid the (nx, ny, d) intermediate.
     return jax.vmap(lambda xi: jnp.linalg.norm(xi - y, ord=p, axis=-1))(x)
 
@@ -160,13 +156,12 @@ def pted_chunk_numpy(
         x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_numpy
     )
     permute_stats = []
+    z = np.concatenate((x, y), axis=0)
     for _ in trange(permutations, disable=not prog_bar):
-        z = np.concatenate((x, y), axis=0)
         z = z[np.random.permutation(len(z))]
-        x, y = z[:nx], z[nx:]
         permute_stats.append(
             _energy_distance_estimate(
-                x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_numpy
+                z[:nx], z[nx:], chunk_size, metric=metric, energy_distance_fn=_energy_distance_numpy
             )
         )
     return test_stat, permute_stats
@@ -190,13 +185,12 @@ def pted_chunk_torch(
         x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_torch
     )
     permute_stats = []
+    z = torch.cat((x, y), dim=0)
     for _ in trange(permutations, disable=not prog_bar):
-        z = torch.cat((x, y), dim=0)
         z = z[torch.randperm(len(z))]
-        x, y = z[:nx], z[nx:]
         permute_stats.append(
             _energy_distance_estimate(
-                x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_torch
+                z[:nx], z[nx:], chunk_size, metric=metric, energy_distance_fn=_energy_distance_torch
             )
         )
     return test_stat, permute_stats
@@ -299,13 +293,12 @@ def pted_chunk_jax(
         x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_jax
     )
     permute_stats = []
+    z = jnp.concatenate([x, y], axis=0)
     for _ in trange(permutations, disable=not prog_bar):
-        z = jnp.concatenate([x, y], axis=0)
         z = z[np.random.permutation(len(z))]
-        x, y = z[:nx], z[nx:]
         permute_stats.append(
             _energy_distance_estimate(
-                x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_jax
+                z[:nx], z[nx:], chunk_size, metric=metric, energy_distance_fn=_energy_distance_jax
             )
         )
     return test_stat, permute_stats
diff --git a/tests/test_pted.py b/tests/test_pted.py
index bb3c925..56853fa 100644
--- a/tests/test_pted.py
+++ b/tests/test_pted.py
@@ -31,6 +31,7 @@ def test_inputs_extra_dims():
     if torch is None:
         pytest.skip("torch not installed")
     # Test with torch tensors
+    torch.manual_seed(42)
     g = torch.randn(100, 30, 30)
     s = torch.randn(50, 100, 30, 30)
     p = pted.pted_coverage_test(g, s)
@@ -84,6 +85,7 @@ def test_pted_torch():
 
 
 def test_pted_coverage_full():
+    np.random.seed(42)
     g = np.random.normal(size=(100, 10))  # ground truth (n_simulations, n_dimensions)
     s = np.random.normal(
         size=(200, 100, 10)
@@ -162,6 +164,7 @@ def test_pted_chunk_size_fallback():
 
 def test_pted_coverage_edgecase():
     # Test with single simulation
+    np.random.seed(42)
     g = np.random.normal(size=(1, 10))
     s = np.random.normal(size=(100, 1, 10))
     p = pted.pted_coverage_test(g, s)
@@ -169,6 +172,7 @@ def test_pted_coverage_edgecase():
 
 
 def test_pted_coverage_progress_bar(capsys):
+    np.random.seed(42)
     g = np.random.normal(size=(42, 10))
     s = np.random.normal(size=(100, 42, 10))
     pted.pted_coverage_test(g, s)
@@ -185,6 +189,7 @@ def test_pted_coverage_progress_bar(capsys):
 def test_pted_coverage_overunder():
     if torch is None:
         pytest.skip("torch not installed")
+    torch.manual_seed(42)
     g = torch.randn(100, 3)
     s = torch.randn(50, 100, 3)
     with pytest.warns(pted.utils.OverconfidenceWarning):
@@ -194,6 +199,7 @@ def test_pted_coverage_overunder():
 
 
 def test_sbc_histogram():
+    np.random.seed(42)
     g = np.random.normal(size=(100, 10))  # ground truth (nsim, ndim)
     s = np.random.normal(size=(150, 100, 10))  # posterior samples (nsamp, nsim, ndim)
 
@@ -202,6 +208,7 @@ def test_sbc_histogram():
 
 
 def test_pit_plot_coverage_test():
+    np.random.seed(42)
     g = np.random.normal(size=(100, 10))  # ground truth (nsim, ndim)
     s = np.random.normal(size=(150, 100, 10))  # posterior samples (nsamp, nsim, ndim)
 
@@ -212,6 +219,7 @@ def test_pit_plot_coverage_test():
 
 def test_pit_plot_utility_direct():
     """pit_plot utility function creates a file and handles edge cases."""
+    np.random.seed(42)
     pvals = np.random.uniform(size=50)
     pted.utils.pit_plot(pvals, "pit_direct.pdf")
     assert os.path.exists("pit_direct.pdf"), "PIT plot file was not created"
@@ -276,6 +284,7 @@ def test_pted_coverage_jax():
     if jax is None:
         pytest.skip("jax not installed")
 
+    np.random.seed(42)
     g = jnp.array(np.random.normal(size=(75, 5)))
     s = jnp.array(np.random.normal(size=(50, 75, 5)))
     p = pted.pted_coverage_test(g, s)
@@ -326,6 +335,7 @@ def test_energy_distance_jax():
     """_energy_distance_jax returns 0 when x and y are identical."""
     if jax is None:
         pytest.skip("jax not installed")
+    np.random.seed(42)
     x = jnp.array(np.random.normal(size=(50, 5)))
     # Identical samples means energy distance should be ~0
     ed = pted.utils._energy_distance_jax(x, x)

From 0b1808def40b4a2ffb8a04976af867a363fb910f Mon Sep 17 00:00:00 2001
From: Connor Stone <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 11:03:41 -0400
Subject: [PATCH 11/12] upgrade versions allowed

---
 .github/workflows/cd.yml | 2 +-
 .github/workflows/ci.yml | 8 ++++----
 pyproject.toml           | 6 +++---
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml
index 6284ed8..21e2290 100644
--- a/.github/workflows/cd.yml
+++ b/.github/workflows/cd.yml
@@ -51,7 +51,7 @@ jobs:
       - uses: actions/setup-python@v5
         name: Install Python
         with:
-          python-version: "3.10"
+          python-version: "3.13"
       - uses: actions/download-artifact@v4
         with:
           name: artifact
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index fa73e2d..1dc49b6 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -22,7 +22,7 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12", "3.13"]
         os: [ubuntu-latest, windows-latest, macOS-latest]
 
     steps:
@@ -51,15 +51,15 @@ jobs:
           python -m pip install --upgrade pip
           pip install pytest pytest-cov wheel
 
-      - name: Install torch # no torch on 3.11 to test no-torch scenario
-        if: ${{ matrix.python-version != '3.11' }}
+      - name: Install torch # no torch on 3.13 to test no-torch scenario
+        if: ${{ matrix.python-version != '3.13' }}
         run: |
           pip install torch
 
       # We only want to install this on one run, because otherwise we'll have
       # duplicate annotations.
       - name: Install error reporter
-        if: ${{ matrix.python-version == '3.10' }}
+        if: ${{ matrix.python-version == '3.13' }}
         run: |
           python -m pip install pytest-github-actions-annotate-failures
 
diff --git a/pyproject.toml b/pyproject.toml
index 0b3aee1..c96f676 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,7 +13,7 @@ authors = [
 ]
 description = "Implementation of a Permutation Test using the Energy Distance for two sample tests and posterior coverage tests"
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.11"
 license = {file = "LICENSE"}
 keywords = [
         "statistics",
@@ -41,14 +41,14 @@ dev = [
     "pytest-cov>=4.1,<5",
     "pytest-mock>=3.12,<4",
     "torch>=2.0,<3",
-    "jax>=0.4,<1",
+    "jax>=0.7,<1",
     "matplotlib",
 ]
 torch = [
     "torch>=2.0,<3",
 ]
 jax = [
-    "jax>=0.4,<1",
+    "jax>=0.7,<1",
 ]
 
 [tool.hatch.metadata.hooks.requirements_txt]

From dea8e0be58f3069bdc2aa668398d13e2db60a40e Mon Sep 17 00:00:00 2001
From: Connor Stone <connorstone628@gmail.com>
Date: Sat, 1 Aug 2026 11:17:50 -0400
Subject: [PATCH 12/12] change random seed since changed torch rnadom behaviour
 caused problem

---
 tests/test_pted.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_pted.py b/tests/test_pted.py
index 56853fa..7d3d166 100644
--- a/tests/test_pted.py
+++ b/tests/test_pted.py
@@ -31,7 +31,7 @@ def test_inputs_extra_dims():
     if torch is None:
         pytest.skip("torch not installed")
     # Test with torch tensors
-    torch.manual_seed(42)
+    torch.manual_seed(43)
     g = torch.randn(100, 30, 30)
     s = torch.randn(50, 100, 30, 30)
     p = pted.pted_coverage_test(g, s)