polystep 0.12.0__tar.gz → 0.13.0__tar.gz

{polystep-0.12.0 → polystep-0.13.0}/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 0.13.0 - 2026-09-22
+
+- Fix KL row mass, nonfinite inputs, evaluator hooks, and chunk handling.
+- Reuse shared MLP forward work and load only requested MNIST images.
+- Fix example checkpoint selection and evaluation budgets; trim demo dependencies.
+
 ## 0.12.0 - 2026-09-17
 
 - Fix cross-entropy targets, stale caches, attention masks, nonfinite best values,

{polystep-0.12.0 → polystep-0.13.0}/CITATION.cff

@@ -6,8 +6,8 @@ authors:
   - family-names: Le
     given-names: An T.
 license: Apache-2.0
-version: 0.12.0
-date-released: "2026-09-17"
+version: 0.13.0
+date-released: "2026-09-22"
 repository-code: "https://github.com/anindex/polystep"
 keywords:
   - gradient-free optimization

{polystep-0.12.0 → polystep-0.13.0}/CONTRIBUTING.md

@@ -26,7 +26,12 @@ Tests use one CPU thread by default; `POLYSTEP_TEST_THREADS` overrides it. Slow
 may download MNIST. Tests importing `experiments/` must skip when that directory is
 absent, since it is excluded from distributions.
 
-CI also builds the source distribution and runs its tests. Before a release:
+Pushes and pull requests test Python 3.11 with PyTorch 2.8 and Python 3.14 with
+current PyTorch. Scheduled and manual runs add Python 3.12/3.13, coverage, real-data
+learning checks, and the shipped tests against the installed wheel. Every run
+builds both distributions and checks the wheel's imports and type marker.
+
+Before a release:
 
 ```bash
 uv build
@@ -39,7 +44,9 @@ Update `polystep.__version__`, the version and date in `CITATION.cff`, and
 `CHANGELOG.md`. `pyproject.toml` reads the package version automatically.
 
 Publishing a GitHub Release tagged `vX.Y.Z` triggers the PyPI workflow. It checks
-version consistency, lint, tests, and package metadata before publishing.
+version consistency, lint, tests, and package metadata before publishing. The
+release also runs the shipped tests against the wheel. Only the separate publish
+job has OIDC permission; it uses the `pypi` environment and PyPI Trusted Publishing.
 
 ## Issues
 

{polystep-0.12.0 → polystep-0.13.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polystep
-Version: 0.12.0
+Version: 0.13.0
 Summary: PyTorch PolyStep Optimizer - gradient-free neural network training via entropic optimal transport
 Author: An T. Le
 License-Expression: Apache-2.0
@@ -27,13 +27,12 @@ License-File: LICENSE
 Requires-Dist: torch>=2.8
 Provides-Extra: examples
 Requires-Dist: numpy>=2.0; extra == "examples"
-Requires-Dist: torchvision>=0.25; extra == "examples"
 Requires-Dist: matplotlib>=3.8; extra == "examples"
 Requires-Dist: Pillow>=10.0; extra == "examples"
 Requires-Dist: gymnasium[classic-control]>=0.29; extra == "examples"
-Requires-Dist: python-sat; extra == "examples"
 Provides-Extra: dev
 Requires-Dist: numpy>=2.0; extra == "dev"
+Requires-Dist: setuptools>=77; extra == "dev"
 Requires-Dist: pytest>=8.0; extra == "dev"
 Requires-Dist: pytest-timeout>=2.2; extra == "dev"
 Requires-Dist: pytest-rerunfailures>=14.0; extra == "dev"
@@ -66,6 +65,8 @@ Dynamic: license-file
 
 Gradient-free neural network training via optimal transport.
 
+Accepted at [Transactions on Machine Learning Research (TMLR)](https://openreview.net/forum?id=8mlcqTTMuU).
+
 PolyStep evaluates finite-radius parameter perturbations and updates parameters
 through a weighted barycenter. It supports spiking networks, quantized layers, hard
 routing, and other models without useful gradients. For differentiable models,
@@ -77,11 +78,9 @@ backpropagation is usually faster.
 
 <p align="center">
   <img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/method_diagram.png" width="840"
-       alt="PolyStep: subspace projection, polytope probes, loss evaluation, and a weighted parameter update.">
+       alt="PolyStep: search coordinates, polytope probes, cost evaluation, soft assignment, and a barycentric update.">
 </p>
 
-> Want to play around with parameters? [**Viet T. Nguyen**](https://vietngth.github.io/) built a gorgeous interactive walkthrough that animates every step of the method -> **[explore the PolyStep visualization](https://vietngth.github.io/polystep-visualization/)**.
-
 ## Installation
 
 ```bash
@@ -178,7 +177,7 @@ studies are reported separately in the [experiment index](experiments/EXPERIMENT
 <!-- BENCH:START -->
 ### Non-differentiable tasks
 
-Test accuracy %. A dash means the cell is not in this release.
+Test accuracy %. A dash means the result is not in this release.
 
 | Task | PolyStep | CMA-ES | OpenAI-ES | SPSA | Adam (surrogate) | Non-diff op |
 |---|---|---|---|---|---|---|
@@ -230,7 +229,8 @@ on the differentiable tasks, and specialized solvers lead on MAX-SAT. See
 
 ## Acknowledgments
 
-A huge thank you to [**Viet**](https://vietngth.github.io/) for building a beautiful interactive [PolyStep visualization](https://vietngth.github.io/polystep-visualization/), it brings the method to life and makes every step click!
+Thanks to [Viet T. Nguyen](https://vietngth.github.io/) for the
+[interactive visualization](https://vietngth.github.io/polystep-visualization/).
 
 ## License
 

{polystep-0.12.0 → polystep-0.13.0}/README.md

@@ -8,6 +8,8 @@
 
 Gradient-free neural network training via optimal transport.
 
+Accepted at [Transactions on Machine Learning Research (TMLR)](https://openreview.net/forum?id=8mlcqTTMuU).
+
 PolyStep evaluates finite-radius parameter perturbations and updates parameters
 through a weighted barycenter. It supports spiking networks, quantized layers, hard
 routing, and other models without useful gradients. For differentiable models,
@@ -19,11 +21,9 @@ backpropagation is usually faster.
 
 <p align="center">
   <img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/method_diagram.png" width="840"
-       alt="PolyStep: subspace projection, polytope probes, loss evaluation, and a weighted parameter update.">
+       alt="PolyStep: search coordinates, polytope probes, cost evaluation, soft assignment, and a barycentric update.">
 </p>
 
-> Want to play around with parameters? [**Viet T. Nguyen**](https://vietngth.github.io/) built a gorgeous interactive walkthrough that animates every step of the method -> **[explore the PolyStep visualization](https://vietngth.github.io/polystep-visualization/)**.
-
 ## Installation
 
 ```bash
@@ -120,7 +120,7 @@ studies are reported separately in the [experiment index](experiments/EXPERIMENT
 <!-- BENCH:START -->
 ### Non-differentiable tasks
 
-Test accuracy %. A dash means the cell is not in this release.
+Test accuracy %. A dash means the result is not in this release.
 
 | Task | PolyStep | CMA-ES | OpenAI-ES | SPSA | Adam (surrogate) | Non-diff op |
 |---|---|---|---|---|---|---|
@@ -172,7 +172,8 @@ on the differentiable tasks, and specialized solvers lead on MAX-SAT. See
 
 ## Acknowledgments
 
-A huge thank you to [**Viet**](https://vietngth.github.io/) for building a beautiful interactive [PolyStep visualization](https://vietngth.github.io/polystep-visualization/), it brings the method to life and makes every step click!
+Thanks to [Viet T. Nguyen](https://vietngth.github.io/) for the
+[interactive visualization](https://vietngth.github.io/polystep-visualization/).
 
 ## License
 

{polystep-0.12.0 → polystep-0.13.0}/docs/performance.md

@@ -32,6 +32,10 @@ Delta paths reuse base activations within each candidate sweep. Reuse applies to
 supported builtin layers and ends before the next batch or fidelity level.
 Custom layers can declare the contracts in the [API reference](api_overview.md#custom-layers).
 
+Dense MLP evaluation also shares activations until the first changed layer when
+unchanged candidate parameters are expanded views. Shared weights use one linear
+operation across candidates. `chunk_size` bounds activation batches on this path.
+
 ## Subspaces and memory
 
 Use `HybridSubspace.from_layout(..., max_subspace_dim=...)` to bound candidate count;
@@ -121,7 +125,7 @@ PYTHONPATH=src:. python experiments/scripts/bench_forward_backends.py --profile
 PYTHONPATH=src:. python experiments/scripts/bench_forward_backends.py --prefix --repeats 7 --warmup 2
 PYTHONPATH=src:. python experiments/scripts/bench_forward_backends.py --attention --batch 16
 PYTHONPATH=src:. python experiments/scripts/bench_forward_backends.py --compile --profile
-PYTHONPATH=src:. python experiments/scripts/bench_polytope.py --suite --seeds 0 1 2 3 4 --budget 4096 --wall
+PYTHONPATH=src:. python experiments/scripts/bench_polytope.py --seeds 0 1 2 3 4 --budget 4096 --wall
 ```
 
 Use `--device cuda` for GPU measurements and `--output` to keep separate result files.

{polystep-0.12.0 → polystep-0.13.0}/docs/reproducibility.md

@@ -32,11 +32,13 @@ python experiments/scripts/tune_gallery.py
 python experiments/scripts/tune_polystep.py
 bash experiments/scripts/run_revision.sh
 python experiments/scripts/check_grid_complete.py experiments/results/revision
-python experiments/scripts/generate_paper_tables.py \
+python experiments/scripts/generate_tables.py \
     --results-dir experiments/results/revision \
     --readme README.md --index experiments/EXPERIMENT_INDEX.md
 ```
 
+The table generator updates Markdown only and rejects test-selected results.
+
 Individual runners accept `--methods` and `--seeds`:
 
 ```bash
@@ -73,3 +75,6 @@ python experiments/scripts/aggregate_results.py experiments/results/revision --w
 
 See [determinism](determinism.md), [controlled experiments](../experiments/CONTROLLED_EXPERIMENTS.md),
 and [forward benchmarks](performance.md#benchmarks) for their specific checks.
+
+With the datasets and CUDA available, run the practical-study update checks with
+`PYTHONPATH=src:. pytest tests/test_practical.py -m gpu`.

{polystep-0.12.0 → polystep-0.13.0}/pyproject.toml

@@ -38,15 +38,14 @@ Issues = "https://github.com/anindex/polystep/issues"
 # Lightweight: just enough for examples/ demos
 examples = [
     "numpy>=2.0",
-    "torchvision>=0.25",
     "matplotlib>=3.8",
     "Pillow>=10.0",
     "gymnasium[classic-control]>=0.29",
-    "python-sat",  # example 04 generates its 3-SAT instance through pysat
 ]
 # Core dev: lint + test the library itself (no experiment deps)
 dev = [
     "numpy>=2.0",  # the test suite loads MNIST through it
+    "setuptools>=77",  # PyTorch 2.8's compiler imports it at runtime
     "pytest>=8.0",
     "pytest-timeout>=2.2",
     "pytest-rerunfailures>=14.0",

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/__init__.py

@@ -1,6 +1,6 @@
 """polystep: PyTorch PolyStep Optimizer."""
 
-__version__ = "0.12.0"
+__version__ = "0.13.0"
 
 __all__ = [
     # Core solver

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/ask_tell.py

@@ -11,6 +11,7 @@ import torch
 from ._compiled import _barycentric_projection, _rotate_and_translate
 from .geometry import get_orthoplex_vertices, get_random_rotation_matrices
 from .solvers import SinkhornSolver, SoftmaxSolver
+from .solvers._shared import validate_positive
 
 __all__ = ["PolyStepES", "minimize"]
 
@@ -52,8 +53,7 @@ class PolyStepES:
             raise ValueError(f"dim must be >= 1, got {dim}.")
         if num_particles < 1:
             raise ValueError(f"num_particles must be >= 1, got {num_particles}.")
-        if not epsilon > 0:
-            raise ValueError(f"epsilon must be > 0, got {epsilon}.")
+        validate_positive(epsilon, "epsilon")
         if not (math.isfinite(step_radius) and step_radius >= 0):
             raise ValueError(f"step_radius must be finite and >= 0, got {step_radius}.")
         self.dim = dim

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/benchmarks/utils.py

@@ -64,11 +64,12 @@ def _download_mnist(data_dir: str) -> None:
         download_file(MNIST_URL + filename, filepath)
 
 
-def _load_mnist_images(filepath: str) -> np.ndarray:
+def _load_mnist_images(filepath: str, limit: int = 0) -> np.ndarray:
     """Load MNIST images from gzipped IDX file."""
     with gzip.open(filepath, "rb") as f:
         _magic, num, rows, cols = pystruct.unpack(">IIII", f.read(16))
-        images = np.frombuffer(f.read(), dtype=np.uint8)
+        num = min(num, limit) if limit > 0 else num
+        images = np.frombuffer(f.read(num * rows * cols), dtype=np.uint8)
         images = images.reshape(num, 1, rows, cols)
     return images.astype(np.float32) / 255.0
 
@@ -92,9 +93,9 @@ def get_mnist_loaders(
     data_dir = data_dir or _default_data_dir("mnist")
     _download_mnist(data_dir)
 
-    train_images = _load_mnist_images(os.path.join(data_dir, MNIST_FILES["train_images"]))
+    train_images = _load_mnist_images(os.path.join(data_dir, MNIST_FILES["train_images"]), max_train)
     train_labels = _load_mnist_labels(os.path.join(data_dir, MNIST_FILES["train_labels"]))
-    test_images = _load_mnist_images(os.path.join(data_dir, MNIST_FILES["test_images"]))
+    test_images = _load_mnist_images(os.path.join(data_dir, MNIST_FILES["test_images"]), max_test)
     test_labels = _load_mnist_labels(os.path.join(data_dir, MNIST_FILES["test_labels"]))
 
     if normalize:

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/cost_nn.py

@@ -8,6 +8,7 @@ the one parameter a candidate perturbs.
 from __future__ import annotations
 
 import warnings
+from bisect import bisect_right
 from contextlib import contextmanager
 from typing import Callable, Optional, TYPE_CHECKING, Union
 
@@ -113,7 +114,7 @@ def _batched_loss_kind(loss_fn) -> Optional[str]:
     A configured loss has no reproduction and falls back to vmap; a subclass that
     overrides ``forward`` would score candidates on a different objective.
     """
-    if getattr(loss_fn, "reduction", None) != "mean":
+    if getattr(loss_fn, "reduction", None) != "mean" or _has_hooks(loss_fn):
         return None
     for base, kind in ((nn.CrossEntropyLoss, "cross_entropy"), (nn.MSELoss, "mse"), (nn.L1Loss, "l1")):
         if not isinstance(loss_fn, base):
@@ -208,6 +209,8 @@ class NNCostEvaluator:
         per_sample: bool = False,
         autocast_dtype: Optional[torch.dtype] = None,
     ):
+        if chunk_size not in (None, "auto") and (type(chunk_size) is not int or chunk_size < 1):
+            raise ValueError("chunk_size must be None, 'auto', or a positive integer")
         self.model = model
         self.loss_fn = loss_fn
         self.autocast_dtype = autocast_dtype
@@ -297,6 +300,7 @@ class NNCostEvaluator:
         self._subspace_backup = None
         self._cast_cache = None
         self._input_dtype_cache = _UNSET
+        self._chunk_size_cached = _UNSET
         loss_kind = _batched_loss_kind(self.loss_fn)
         self._batched_linear = (
             BatchedLinearEvaluator.try_build(self.model, self.loss_fn, loss_kind) if loss_kind is not None else None
@@ -385,7 +389,18 @@ class NNCostEvaluator:
             and not self.per_sample
             and not (self._batched_linear.loss_kind == "cross_entropy" and targets.dim() != 1)
         ):
-            return self._batched_linear.evaluate(stacked_params, inputs, targets)
+            n = next(iter(stacked_params.values())).shape[0]
+            chunk = self.chunk_size
+            if chunk is None or n <= chunk:
+                return self._batched_linear.evaluate(stacked_params, inputs, targets)
+            return torch.cat(
+                [
+                    self._batched_linear.evaluate(
+                        {k: v[start : start + chunk] for k, v in stacked_params.items()}, inputs, targets
+                    )
+                    for start in range(0, n, chunk)
+                ]
+            )
 
         if self._vmap_failed:
             result = self._evaluate_loop(stacked_params, inputs, targets)
@@ -856,7 +871,7 @@ def _plan_children(model: nn.Module):
     for name, mod in model._modules.items():
         if mod is None:
             continue
-        if isinstance(mod, nn.Sequential):
+        if isinstance(mod, nn.Sequential) and type(mod).forward is nn.Sequential.forward and not _has_hooks(mod):
             for sub, m in mod._modules.items():
                 if m is not None:
                     yield f"{name}.{sub}", m
@@ -995,42 +1010,49 @@ class BatchedLinearEvaluator:
     ) -> torch.Tensor:
         """Batched forward via bmm for Linear layers."""
         N = next(iter(stacked_params.values())).shape[0]
-        # Expand the input across the N param configs.
-        if inputs.dim() == 2:
-            x = inputs.unsqueeze(0).expand(N, -1, -1)  # (N, B, in_feat)
-        else:
-            # Flatten spatial dims for non-2D inputs.
-            x = inputs.reshape(inputs.shape[0], -1).unsqueeze(0).expand(N, -1, -1)
+        # Keep the prefix shared until a candidate actually changes a layer.
+        x = inputs if inputs.dim() == 2 else inputs.reshape(inputs.shape[0], -1)
 
         for name, tag, module in self._layer_keys:
             if tag == "linear":
                 w_key = f"{name}.weight"
                 b_key = f"{name}.bias"
                 wq, bq = _weight_transforms(module)
-                # Frozen params are constant across candidates, so the module's tensor
-                # stands in; bmm does not broadcast its batch dim, hence the expand.
                 W = stacked_params.get(w_key)
                 if W is None:
-                    W = module.weight.unsqueeze(0).expand(N, -1, -1)
+                    W = module.weight
+                elif W.stride(0) == 0:
+                    W = W[0]  # expanded view: every candidate shares this weight
                 if wq is not None:
                     # Elementwise, so transforming the whole stack equals per-candidate.
                     W = wq(W)
                 bias = stacked_params.get(b_key)
                 if bias is None and getattr(module, "bias", None) is not None:
-                    bias = module.bias.unsqueeze(0)
+                    bias = module.bias
+                elif bias is not None and bias.stride(0) == 0:
+                    bias = bias[0]
                 if bias is not None and bq is not None:
                     bias = bq(bias)
-                # baddbmm folds the bias into the same kernel.
-                if bias is not None:
-                    x = torch.baddbmm(bias.unsqueeze(1), x, W.transpose(1, 2))
+                if W.dim() == 2:
+                    # One GEMM for shared weights, including a shared prefix.
+                    x = nn.functional.linear(x, W, bias if bias is not None and bias.dim() == 1 else None)
+                    if bias is not None and bias.dim() == 2:
+                        x = x + bias.unsqueeze(1)
                 else:
-                    x = torch.bmm(x, W.transpose(1, 2))
+                    if x.dim() == 2:
+                        x = x.unsqueeze(0).expand(N, -1, -1)
+                    if bias is not None:
+                        x = torch.baddbmm(bias.unsqueeze(-2), x, W.transpose(1, 2))
+                    else:
+                        x = torch.bmm(x, W.transpose(1, 2))
             elif tag in ("flatten", "dropout"):
                 pass  # input is pre-flattened; dropout is identity in eval mode
             else:
                 # Apply the real module so its configuration is exact.
                 x = module(x)
 
+        if x.dim() == 2:
+            x = x.unsqueeze(0).expand(N, -1, -1)
         return _reduce_per_candidate(x, targets, self.loss_fn, self.loss_kind)
 
 
@@ -1046,7 +1068,7 @@ def _resolve_layout_entry(self, offsets: torch.Tensor, pdim: int, span=None):
     else:
         lo, hi = span
 
-    idx = int(torch.searchsorted(self._starts, torch.tensor(lo), right=True)) - 1
+    idx = bisect_right(self._starts, lo) - 1
     if idx < 0 or idx >= len(self._keys):
         return None
     entry = self._entry_by_key[self._keys[idx]]
@@ -1078,7 +1100,7 @@ class SiteVmapEvaluator:
         self._buffers = buffers
         self._owner = owner
         self._entry_by_key = {e.key: e for e in layout.entries}
-        self._starts = torch.tensor([e.offset for e in layout.entries], dtype=torch.long)
+        self._starts = [e.offset for e in layout.entries]
         self._keys = [e.key for e in layout.entries]
         self._compiled_sites = {}
         self._compile_failed_sites = set()
@@ -1245,7 +1267,7 @@ class SparseDeltaEvaluator:
                 self._site[f"{name}.bias"] = (idx, False)
         self._entry_by_key = {e.key: e for e in layout.entries}
         # Flat spans of the entries this path can perturb, for the containment test.
-        self._starts = torch.tensor([e.offset for e in layout.entries], dtype=torch.long)
+        self._starts = [e.offset for e in layout.entries]
         self._keys = [e.key for e in layout.entries]
 
     @classmethod

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/optimizer.py

@@ -65,6 +65,7 @@ from .solvers import (
     TopKMeanSolver,
 )
 from .solver import SolverState
+from .solvers._shared import validate_positive
 from .adaptive_subspace import AdaptiveSubspace
 from .cma import compute_cma_hyperparameters
 from .cma_subspace import CMAAdaptiveSubspace
@@ -1112,20 +1113,13 @@ class PolyStepOptimizer(SerializationMixin):
         """Actual projection type in use ('dense' or 'sparse')."""
         return self._actual_projection_type
 
-    @staticmethod
-    def _check_temperature(value: float, name: str) -> float:
-        """Reject a non-positive temperature the fused kernel would divide by unguarded."""
-        if not value > 0:
-            raise ValueError(f"{name} must resolve to > 0, got {value}.")
-        return value
-
     def _get_epsilon(self, iteration: int) -> float:
         """Resolve epsilon at current iteration."""
         if self._progressive_epsilon is not None:
-            return self._check_temperature(self._progressive_epsilon.at(iteration), "epsilon")
+            return validate_positive(self._progressive_epsilon.at(iteration), "epsilon")
         if hasattr(self.epsilon, "at"):
-            return self._check_temperature(self.epsilon.at(iteration), "epsilon")
-        return self.epsilon
+            return validate_positive(self.epsilon.at(iteration), "epsilon")
+        return validate_positive(self.epsilon, "epsilon")
 
     def _get_step_radius(self, iteration: int) -> float:
         """Resolve step_radius at current iteration (supports schedule objects)."""
@@ -1220,8 +1214,8 @@ class PolyStepOptimizer(SerializationMixin):
         if self.ent_epsilon is None:
             return None
         if hasattr(self.ent_epsilon, "at"):
-            return self._check_temperature(self.ent_epsilon.at(iteration), "ent_epsilon")
-        return self.ent_epsilon
+            return validate_positive(self.ent_epsilon.at(iteration), "ent_epsilon")
+        return validate_positive(self.ent_epsilon, "ent_epsilon")
 
     def _bf16_supported(self) -> bool:
         """Whether BF16 is supported on the model's device."""

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/solver.py

@@ -12,7 +12,7 @@ from .costs import compute_cost_matrix
 from .epsilon import feed_solver_stats, resolve_radius, LinearEpsilon
 from .geometry import get_random_rotation_matrices, POLYTOPE_MAP
 from .solvers import SinkhornSolver, SoftmaxSolver
-from .solvers._shared import solver_health
+from .solvers._shared import solver_health, validate_positive
 
 
 @dataclass
@@ -173,16 +173,16 @@ class PolyStep:
     def _get_epsilon(self, iteration: int) -> float:
         """Resolve epsilon at ``iteration``; duck-types on ``.at()`` so schedules and plain floats both work."""
         if hasattr(self.epsilon, "at"):
-            return self.epsilon.at(iteration)
-        return self.epsilon
+            return validate_positive(self.epsilon.at(iteration), "epsilon")
+        return validate_positive(self.epsilon, "epsilon")
 
     def _get_ent_epsilon(self, iteration: int) -> Optional[float]:
         """Resolve ent_epsilon at current iteration (supports schedule objects)."""
         if self.ent_epsilon is None:
             return None
         if hasattr(self.ent_epsilon, "at"):
-            return self.ent_epsilon.at(iteration)
-        return self.ent_epsilon
+            return validate_positive(self.ent_epsilon.at(iteration), "ent_epsilon")
+        return validate_positive(self.ent_epsilon, "ent_epsilon")
 
     @torch.inference_mode()
     def step(

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/solvers/_shared.py

@@ -10,13 +10,14 @@ import torch
 from .base import SolverResult
 
 
-def validate_positive(value: float, name: str, context: str = "") -> None:
-    """Raise ValueError unless ``value > 0``."""
-    if not value > 0:
-        msg = f"{name} must be > 0, got {value}."
+def validate_positive(value: float, name: str, context: str = "") -> float:
+    """Return a positive finite value, or raise ValueError."""
+    if not (math.isfinite(value) and value > 0):
+        msg = f"{name} must be > 0 and finite, got {value}."
         if context:
             msg += " " + context
         raise ValueError(msg)
+    return value
 
 
 def solver_health(transport: torch.Tensor, displacement: torch.Tensor, step_radius: float):

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/solvers/greedy.py

@@ -89,13 +89,15 @@ class TopKMeanSolver:
         _, topk_indices = C_raw.topk(k_eff, dim=-1, largest=False)  # (P, k_eff)
 
         # topk does not honour the sanitize mask, so drop infeasible picks and spread mass over what is left; the row still sums to a.
-        feasible = ~(torch.isnan(cost_matrix) | (cost_matrix == float("inf")))
+        feasible = torch.isfinite(cost_matrix)
         keep = feasible.gather(1, topk_indices)
         keep = keep | ~keep.any(dim=1, keepdim=True)
 
         transport = torch.zeros(P, V, device=device, dtype=dtype)
         mass = a.unsqueeze(1) / keep.sum(dim=1, keepdim=True) * keep
         transport.scatter_(1, topk_indices, mass.to(dtype))
+        informative = (C_raw.amax(dim=-1) > C_raw.amin(dim=-1)).unsqueeze(1)
+        transport = torch.where(informative, transport, a.unsqueeze(1).expand(P, V) / V)
 
         # scale_cost is unused: topk is invariant to a positive divisor.
         ent_cost = (C_raw * transport).sum().item()

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/solvers/kl_softmax.py

@@ -11,7 +11,6 @@ import torch
 from ._shared import (
     align_dual,
     align_marginal,
-    exp_plan,
     prepare_cost,
     validate_positive,
     warn_tiny_temperature,
@@ -33,10 +32,7 @@ class KLSoftmaxSolver:
     last_marginal_violation: Optional[float] = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
-        if self.epsilon <= 0:
-            raise ValueError(
-                f"epsilon must be > 0, got {self.epsilon}. epsilon is the entropic temperature and must be positive."
-            )
+        validate_positive(self.epsilon, "epsilon")
         if self.lam < 0:
             raise ValueError(f"lam must be >= 0, got {self.lam}. lam is the KL penalty on the column marginal.")
         if self.threshold < 0:
@@ -51,7 +47,10 @@ class KLSoftmaxSolver:
             return 0.0
         if math.isinf(self.lam):
             return 1.0
-        return float(self.lam / (self.lam + self.epsilon))
+        if self.lam < self.epsilon:
+            ratio = self.lam / self.epsilon
+            return float(ratio / (1.0 + ratio))
+        return float(1.0 / (1.0 + self.epsilon / self.lam))
 
     def solve(
         self,
@@ -98,7 +97,6 @@ class KLSoftmaxSolver:
 
             # Softmax limit, closed form in one iteration. The bound is the dtype's smallest normal: below it the /alpha residual becomes 0/0 = NaN.
             if alpha < torch.finfo(dtype).tiny:
-                f = eps * (log_a - torch.logsumexp(-C / eps, dim=1))
                 g = torch.zeros_like(g)
                 converged = True
                 n_iters = 1
@@ -129,11 +127,13 @@ class KLSoftmaxSolver:
                             break
                     f, g = f_new, g_new
 
-                # One more f-update so P1 == a holds, since the loop leaves f one step behind g.
-                f = eps * (log_a - torch.logsumexp((g.unsqueeze(0) - C) / eps, dim=1))
-
-            # Clamped exponent, not a zero-fill: zeroing overflowed entries would drop the mass that makes P1 == a.
-            P = exp_plan(f, g, C, eps)
+            # Enforce rows directly: exp((f + g - C)/eps) loses their mass when
+            # f and C nearly cancel. Shift before division, including at lam=0.
+            scores = g.unsqueeze(0) - C
+            row_max = scores.amax(dim=1, keepdim=True)
+            logits = (scores - row_max) / eps
+            P = torch.softmax(logits, dim=1) * a.unsqueeze(1)
+            f = eps * (log_a - torch.logsumexp(logits, dim=1)) - row_max.squeeze(1)
 
             # Undo both frame changes so cost is <C_raw, P> (sum(P) == a.sum()).
             cost = ((C * P).sum() * cost_scale + cost_shift * a.sum()).item()

{polystep-0.12.0 → polystep-0.13.0}/src/polystep/solvers/sinkhorn.py

@@ -76,11 +76,7 @@ class SinkhornSolver:
 
     def __post_init__(self):
         """Initialize compiled function registry and validate parameters."""
-        if self.epsilon <= 0:
-            raise ValueError(
-                f"epsilon must be > 0, got {self.epsilon}. "
-                f"A zero or negative epsilon causes division by zero in log-domain Sinkhorn iterations."
-            )
+        validate_positive(self.epsilon, "epsilon")
         if self.omega < 0.5 or self.omega > 1.95:
             raise ValueError(
                 f"omega must be in [0.5, 1.95], got {self.omega}. "