polystep 0.7.0__tar.gz → 0.8.0__tar.gz

{polystep-0.7.0 → polystep-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polystep
-Version: 0.7.0
+Version: 0.8.0
 Summary: PyTorch PolyStep Optimizer - gradient-free neural network training via entropic optimal transport
 Author: An T. Le
 License-Expression: Apache-2.0
@@ -31,7 +31,7 @@ Requires-Dist: pytest>=8.0; extra == "dev"
 Requires-Dist: pytest-timeout>=2.2; extra == "dev"
 Requires-Dist: pytest-rerunfailures>=14.0; extra == "dev"
 Requires-Dist: pytest-cov>=5.0; extra == "dev"
-Requires-Dist: ruff>=0.15.12; extra == "dev"
+Requires-Dist: ruff==0.15.12; extra == "dev"
 Provides-Extra: experiments
 Requires-Dist: scipy>=1.11; extra == "experiments"
 Requires-Dist: statsmodels>=0.14; extra == "experiments"
@@ -39,6 +39,11 @@ Requires-Dist: pandas>=2.0; extra == "experiments"
 Requires-Dist: python-sat; extra == "experiments"
 Requires-Dist: torchvision>=0.25; extra == "experiments"
 Requires-Dist: cma>=3.3; extra == "experiments"
+Requires-Dist: evotorch>=0.5; extra == "experiments"
+Requires-Dist: nevergrad>=1.0; extra == "experiments"
+Requires-Dist: snntorch>=0.7; extra == "experiments"
+Requires-Dist: tabulate>=0.9; extra == "experiments"
+Requires-Dist: datasets>=2.14; extra == "experiments"
 Requires-Dist: gymnasium>=0.29; extra == "experiments"
 Provides-Extra: rl
 Requires-Dist: gymnasium[box2d]>=0.29; extra == "rl"
@@ -68,7 +73,7 @@ Based on the Sinkhorn Step algorithm ([Le et al., NeurIPS 2023](https://arxiv.or
 4. **Update** the parameters by barycentric projection from the weighted vertices.
 
 <p align="center">
-  <img src="docs/figures/method_diagram.png" width="840"
+  <img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/method_diagram.png" width="840"
        alt="One PolyStep: subspace projection, polytope probes, cost matrix, soft entropic-OT assignment, and barycentric projection, plus the softmax-to-full-OT solver continuum.">
 </p>
 
@@ -161,8 +166,8 @@ See [`examples/`](examples/) for runnable demos covering SNN, RL, MAX-SAT, MNIST
 
 <table>
   <tr>
-    <td align="center"><img src="docs/figures/polystep_snn_progress.gif" width="330"><br><sub>Spiking net (hard LIF thresholds) trained with forward passes only</sub></td>
-    <td align="center"><img src="docs/figures/rl_cartpole_policy.gif" width="330"><br><sub>CartPole policy search: no value function, no gradients</sub></td>
+    <td align="center"><img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/polystep_snn_progress.gif" width="330"><br><sub>Spiking net (hard LIF thresholds) trained with forward passes only</sub></td>
+    <td align="center"><img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/rl_cartpole_policy.gif" width="330"><br><sub>CartPole policy search: no value function, no gradients</sub></td>
   </tr>
 </table>
 
@@ -182,7 +187,7 @@ If your model is fully differentiable, Adam/SGD will be faster and more accurate
 
 Zeroth-order optimization splits into two camps. Memory-efficient fine-tuning (MeZO and successors) runs zeroth-order updates on differentiable networks to skip the backprop tape, but still assumes a useful local gradient. Evolution strategies and SPSA (CMA-ES, OpenAI-ES) optimize black-box objectives by estimating a local gradient from small perturbations.
 
-PolyStep is a randomized direct search: it probes a finite-radius polytope around the current parameters and moves toward the lowest-cost vertices through an optimal-transport barycenter. When the loss is piecewise-constant or genuinely non-differentiable, the local gradient is zero almost everywhere, so finite-difference estimates carry no signal and ES/SPSA stall, while a finite radius steps across the flat regions. This is the setting where direct search (generalized pattern search, MADS) has convergence guarantees on nonsmooth and discontinuous objectives that gradient-estimate methods lack. Example 09 shows the separation on a hard decision tree.
+PolyStep is a randomized direct search: it probes a finite-radius polytope around the current parameters and moves toward the lowest-cost vertices through an optimal-transport barycenter. When the loss is piecewise-constant or truly non-differentiable, the local gradient is zero almost everywhere, so finite-difference estimates carry no signal and ES/SPSA stall, while a finite radius steps across the flat regions. This is the setting where direct search (generalized pattern search, MADS) has convergence guarantees on nonsmooth and discontinuous objectives that gradient-estimate methods lack. Example 09 shows the separation on a hard decision tree.
 
 Recent gradient-free work targets the same regime, including low-rank evolution strategies for spiking networks (arXiv:2605.30361) and gradient-free trust regions for recurrent spiking networks (arXiv:2601.21572). PolyStep is complementary: a general-purpose direct-search optimizer with subspace compression.
 
@@ -227,7 +232,9 @@ Among gradient-free optimizers PolyStep is the strongest on every task here, at
 - **Softmax OT solver** with an entropic Sinkhorn alternative.
 - **Subspace compression**: `HybridSubspace` (recommended), `AdaptiveSubspace`, and sparse projection for very large models.
 - **Block-wise OT** for per-layer decomposition.
-- **`torch.compile`** opt-in on hot paths.
+- **`torch.compile`** opt-in on the OT hot paths (`compile=True`) and on the
+  candidate forward (`compile_evaluator` = fusion; `compile_forward` = CUDA
+  graphs on the in-place path). See [`docs/performance.md`](docs/performance.md).
 - **Vmap-safe layers**: drop-in attention and LSTM that play nicely with `torch.vmap`.
 - **Sub-linear memory**: forward-only evaluation, no BPTT activation tape (~30x savings at long SNN horizons).
 - **CMA-ES inspired adaptation** of subspace covariance (experimental, monolithic step only; `use_adaptive_radius` is the stable default).
@@ -259,6 +266,7 @@ docs/                  API overview, reproducibility guide
 | [`examples/`](examples/) | 8 runnable demos with output figures |
 | [`experiments/`](experiments/) | Full paper reproduction harness |
 | [`docs/api_overview.md`](docs/api_overview.md) | API reference |
+| [`docs/performance.md`](docs/performance.md) | Per-step cost & the compile flags (`compile_evaluator`, `compile_forward`) |
 | [`LIMITATIONS.md`](LIMITATIONS.md) | Known limitations |
 | [`CONTRIBUTING.md`](CONTRIBUTING.md) | Contribution guidelines |
 | [`CHANGELOG.md`](CHANGELOG.md) | Release history |

{polystep-0.7.0 → polystep-0.8.0}/README.md

@@ -20,7 +20,7 @@ Based on the Sinkhorn Step algorithm ([Le et al., NeurIPS 2023](https://arxiv.or
 4. **Update** the parameters by barycentric projection from the weighted vertices.
 
 <p align="center">
-  <img src="docs/figures/method_diagram.png" width="840"
+  <img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/method_diagram.png" width="840"
        alt="One PolyStep: subspace projection, polytope probes, cost matrix, soft entropic-OT assignment, and barycentric projection, plus the softmax-to-full-OT solver continuum.">
 </p>
 
@@ -113,8 +113,8 @@ See [`examples/`](examples/) for runnable demos covering SNN, RL, MAX-SAT, MNIST
 
 <table>
   <tr>
-    <td align="center"><img src="docs/figures/polystep_snn_progress.gif" width="330"><br><sub>Spiking net (hard LIF thresholds) trained with forward passes only</sub></td>
-    <td align="center"><img src="docs/figures/rl_cartpole_policy.gif" width="330"><br><sub>CartPole policy search: no value function, no gradients</sub></td>
+    <td align="center"><img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/polystep_snn_progress.gif" width="330"><br><sub>Spiking net (hard LIF thresholds) trained with forward passes only</sub></td>
+    <td align="center"><img src="https://raw.githubusercontent.com/anindex/polystep/main/docs/figures/rl_cartpole_policy.gif" width="330"><br><sub>CartPole policy search: no value function, no gradients</sub></td>
   </tr>
 </table>
 
@@ -134,7 +134,7 @@ If your model is fully differentiable, Adam/SGD will be faster and more accurate
 
 Zeroth-order optimization splits into two camps. Memory-efficient fine-tuning (MeZO and successors) runs zeroth-order updates on differentiable networks to skip the backprop tape, but still assumes a useful local gradient. Evolution strategies and SPSA (CMA-ES, OpenAI-ES) optimize black-box objectives by estimating a local gradient from small perturbations.
 
-PolyStep is a randomized direct search: it probes a finite-radius polytope around the current parameters and moves toward the lowest-cost vertices through an optimal-transport barycenter. When the loss is piecewise-constant or genuinely non-differentiable, the local gradient is zero almost everywhere, so finite-difference estimates carry no signal and ES/SPSA stall, while a finite radius steps across the flat regions. This is the setting where direct search (generalized pattern search, MADS) has convergence guarantees on nonsmooth and discontinuous objectives that gradient-estimate methods lack. Example 09 shows the separation on a hard decision tree.
+PolyStep is a randomized direct search: it probes a finite-radius polytope around the current parameters and moves toward the lowest-cost vertices through an optimal-transport barycenter. When the loss is piecewise-constant or truly non-differentiable, the local gradient is zero almost everywhere, so finite-difference estimates carry no signal and ES/SPSA stall, while a finite radius steps across the flat regions. This is the setting where direct search (generalized pattern search, MADS) has convergence guarantees on nonsmooth and discontinuous objectives that gradient-estimate methods lack. Example 09 shows the separation on a hard decision tree.
 
 Recent gradient-free work targets the same regime, including low-rank evolution strategies for spiking networks (arXiv:2605.30361) and gradient-free trust regions for recurrent spiking networks (arXiv:2601.21572). PolyStep is complementary: a general-purpose direct-search optimizer with subspace compression.
 
@@ -179,7 +179,9 @@ Among gradient-free optimizers PolyStep is the strongest on every task here, at
 - **Softmax OT solver** with an entropic Sinkhorn alternative.
 - **Subspace compression**: `HybridSubspace` (recommended), `AdaptiveSubspace`, and sparse projection for very large models.
 - **Block-wise OT** for per-layer decomposition.
-- **`torch.compile`** opt-in on hot paths.
+- **`torch.compile`** opt-in on the OT hot paths (`compile=True`) and on the
+  candidate forward (`compile_evaluator` = fusion; `compile_forward` = CUDA
+  graphs on the in-place path). See [`docs/performance.md`](docs/performance.md).
 - **Vmap-safe layers**: drop-in attention and LSTM that play nicely with `torch.vmap`.
 - **Sub-linear memory**: forward-only evaluation, no BPTT activation tape (~30x savings at long SNN horizons).
 - **CMA-ES inspired adaptation** of subspace covariance (experimental, monolithic step only; `use_adaptive_radius` is the stable default).
@@ -211,6 +213,7 @@ docs/                  API overview, reproducibility guide
 | [`examples/`](examples/) | 8 runnable demos with output figures |
 | [`experiments/`](experiments/) | Full paper reproduction harness |
 | [`docs/api_overview.md`](docs/api_overview.md) | API reference |
+| [`docs/performance.md`](docs/performance.md) | Per-step cost & the compile flags (`compile_evaluator`, `compile_forward`) |
 | [`LIMITATIONS.md`](LIMITATIONS.md) | Known limitations |
 | [`CONTRIBUTING.md`](CONTRIBUTING.md) | Contribution guidelines |
 | [`CHANGELOG.md`](CHANGELOG.md) | Release history |

{polystep-0.7.0 → polystep-0.8.0}/pyproject.toml

@@ -43,7 +43,7 @@ dev = [
     "pytest-timeout>=2.2",
     "pytest-rerunfailures>=14.0",
     "pytest-cov>=5.0",
-    "ruff>=0.15.12",
+    "ruff==0.15.12",  # pinned to match CI so `ruff format` never disagrees local vs CI
 ]
 # Full experiment reproduction (paper results, ablations, baselines)
 experiments = [
@@ -52,7 +52,12 @@ experiments = [
     "pandas>=2.0",
     "python-sat",  # unusual pre-release versioning, left unpinned to stay installable
     "torchvision>=0.25",
-    "cma>=3.3",
+    "cma>=3.3",  # run_elevation / run_moe / run_maxsat baselines
+    "evotorch>=0.5",  # CMAES baseline (benchmarks.baselines.train_cmaes)
+    "nevergrad>=1.0",  # nevergrad baseline (benchmarks.baselines)
+    "snntorch>=0.7",  # SNN spike datasets (benchmarks.utils)
+    "tabulate>=0.9",  # results tables (benchmarks.utils)
+    "datasets>=2.14",  # HuggingFace SST-2 / text loaders (benchmarks.utils)
     "gymnasium>=0.29",
 ]
 # RL experiments (Box2D envs)

{polystep-0.7.0 → polystep-0.8.0}/src/polystep/__init__.py

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

{polystep-0.7.0 → polystep-0.8.0}/src/polystep/_step_blockwise.py

@@ -55,9 +55,7 @@ def step_blockwise(opt, closure: Callable) -> float:
     # Probe-radius jitter (Thm. 4.2 condition (iv); no-op when probe_radius_jitter == 0).
     probe_r = opt._apply_probe_radius_jitter(probe_r)
 
-    # Ensure 2D
-    if X.dim() == 1:
-        X = X.unsqueeze(0)
+    # state.X is always 2-D (num_particles, particle_dim); no reshape needed.
 
     # Convert layout-indexed flat to block-indexed flat before splitting.
     # Per-layer blocks pad each entry independently, creating different
@@ -484,9 +482,7 @@ def step_subspace_blockwise(opt, closure: Callable) -> float:
     # Probe-radius jitter (Thm. 4.2 condition (iv); no-op when probe_radius_jitter == 0).
     probe_r = opt._apply_probe_radius_jitter(probe_r)
 
-    # Ensure 2D
-    if X.dim() == 1:
-        X = X.unsqueeze(0)
+    # state.X is always 2-D (num_sub_particles, subspace_particle_dim); no reshape needed.
 
     # Get subspace dimension
     sub_dim = opt.subspace.subspace_dim
@@ -645,7 +641,7 @@ def step_subspace_blockwise(opt, closure: Callable) -> float:
             # projection: (full_dim, sub_dim)
             # reconstruct_batch needs projection argument for AdaptiveSubspace
             # Match dtype with projection for mixed precision compatibility
-            if opt._mixed_precision and state.projection is not None:
+            if opt._mixed_precision and state.projection is not None and state.projection.dtype is not None:
                 base_batch = base_batch.to(dtype=state.projection.dtype)
             if opt._adaptive or opt._cma_subspace:
                 chunk_params = state.subspace.reconstruct_batch(

{polystep-0.7.0 → polystep-0.8.0}/src/polystep/_step_monolithic.py

@@ -474,8 +474,12 @@ def step_monolithic(opt, closure: Callable) -> float:
         ot_result = opt.solver.solve(**solve_kwargs)
     state.last_solve_eps = ot_epsilon
 
-    # Auto-epsilon feedback: update progressive epsilon from solver stats
-    if opt._progressive_epsilon is not None:
+    # Update progressive epsilon from solver stats. Skip in fixed-iteration
+    # Sinkhorn mode (threshold <= 0): there the solver always returns
+    # converged=True with n_iters == max_iterations, so the ratio is 1.0 and the
+    # increase branch would push epsilon to max_epsilon every step. Progressive
+    # epsilon is only paired with SinkhornSolver, so threshold always exists.
+    if opt._progressive_epsilon is not None and getattr(opt.solver, "threshold", 1.0) > 0:
         opt._progressive_epsilon.update(
             n_iters=ot_result.n_iters,
             max_iterations=getattr(opt.solver, "max_iterations", 1),
@@ -972,11 +976,27 @@ def step_monolithic(opt, closure: Callable) -> float:
             state.base_params = new_base
             # Reset subspace coordinates to zero
             state.X = torch.zeros_like(state.X)
-            # Regenerate ALL per-layer projections
-            state.hybrid_projections = hybrid_sub.init_projections(
-                state.X.device,
-                state.X.dtype,
-            )
+            # Regenerate ALL per-layer projections. Default: fresh RANDOM redraw. Opt-in
+            # absorb_aligned_active: bias the new basis toward the just-completed window's productive
+            # directions (displacement-SVD) BEFORE the history is zeroed below. Safe here because the
+            # duals are reset in this same block, so no live Sinkhorn dual is disrupted (the reason
+            # per-step rotation is disabled). Lets a smaller rank track descent across absorbs.
+            if getattr(hybrid_sub, "absorb_aligned_active", False) and state.displacement_history_count > 0:
+                hist = state.displacement_history[: state.displacement_history_count]
+                svd_ratio = hybrid_sub.get_svd_ratio(state.iteration_count, opt.max_iterations or 1)
+                state.hybrid_projections = hybrid_sub._rotate_all_displacement(
+                    state.hybrid_projections,
+                    hist,
+                    svd_ratio,
+                    state.X.device,
+                    state.X.dtype,
+                    state.iteration_count,
+                )
+            else:
+                state.hybrid_projections = hybrid_sub.init_projections(
+                    state.X.device,
+                    state.X.dtype,
+                )
             # Reset displacement history
             state.displacement_history.zero_()
             state.displacement_history_idx = 0

{polystep-0.7.0 → polystep-0.8.0}/src/polystep/adaptive_subspace.py

@@ -542,7 +542,7 @@ class AdaptiveSubspace:
         Returns:
             SVD ratio in [svd_ratio_init, svd_ratio_final].
         """
-        progress = min(1.0, step / max(1, total_steps))
+        progress = min(1.0, step / max(1, total_steps or 1))
         return self.svd_ratio_init + progress * (self.svd_ratio_final - self.svd_ratio_init)
 
     # ------------------------------------------------------------------

{polystep-0.7.0 → polystep-0.8.0}/src/polystep/api.py

@@ -8,6 +8,7 @@ epoch iteration, or batch handling.
 
 from __future__ import annotations
 
+import copy
 from dataclasses import dataclass
 from typing import Callable, List, Optional
 
@@ -76,12 +77,20 @@ class TrainConfig:
         log_every: Step interval for built-in logging. Must be > 0.
         callbacks: List of ``TrainCallback`` instances. ``None`` is
             normalized to an empty list.
+        restore_best: If ``True`` (default), snapshot the model weights
+            whenever the tracked loss reaches a new minimum and restore that
+            best snapshot before ``train()`` returns. Gradient-free search is
+            noisy and updates the model in place, so the *last* step is not
+            necessarily the best; this reverts a late divergence. The tracked
+            loss is the exact full-batch loss when a callback consumes per-step
+            metrics, otherwise the (cheap, already-computed) OT cost proxy.
     """
 
     epochs: int = 10
     batch_size: int = 32
     log_every: int = 10
     callbacks: Optional[List[TrainCallback]] = None
+    restore_best: bool = True
 
     def __post_init__(self):
         if self.epochs <= 0:
@@ -134,7 +143,8 @@ def train(
     evaluator = NNCostEvaluator(
         model,
         loss_fn=loss_fn,
-        compile_vmap=getattr(optimizer, "_compile_evaluator", False),
+        compile_vmap=getattr(optimizer, "compile_evaluator", False),
+        compile_forward=getattr(optimizer, "compile_forward", False),
     )
     callbacks = list(config.callbacks)
     global_step = 0
@@ -146,6 +156,10 @@ def train(
     except StopIteration:
         raise ValueError("Model has no trainable parameters")
 
+    # Track the lowest-loss weights so a noisy run is not left at a late step.
+    best_loss = float("inf")
+    best_state = None
+
     for epoch in range(config.epochs):
         epoch_loss_sum = 0.0
         epoch_loss_count = 0
@@ -158,8 +172,17 @@ def train(
             cost_bs = getattr(optimizer, "cost_batch_size", None)
             if cost_bs is not None and cost_bs < inputs.shape[0]:
                 gen = getattr(optimizer, "_generator", None)
-                if gen is not None and gen.device.type == inputs.device.type:
-                    idx = torch.randperm(inputs.shape[0], device=inputs.device, generator=gen)[:cost_bs]
+                if gen is not None:
+                    if gen.device == inputs.device:
+                        idx = torch.randperm(inputs.shape[0], device=inputs.device, generator=gen)[:cost_bs]
+                    else:
+                        # Devices differ (seeded CPU generator with CUDA inputs,
+                        # or cuda:0 vs cuda:1). randperm needs the generator and
+                        # tensor on the same device: sample on the generator's
+                        # device to keep the seed, then move the index.
+                        idx = torch.randperm(inputs.shape[0], device=gen.device, generator=gen)[:cost_bs].to(
+                            inputs.device
+                        )
                 else:
                     idx = torch.randperm(inputs.shape[0], device=inputs.device)[:cost_bs]
                 cost_inputs = inputs[idx]
@@ -174,32 +197,43 @@ def train(
 
             optimizer.step(closure)
 
-            # Compute training loss separately (not the OT cost)
-            with torch.no_grad():
-                output = model(inputs)
-                train_loss = loss_fn(output, targets).item()
+            state = optimizer.state
+            ot_cost = state.costs[-1] if state.costs else 0.0
+
+            if callbacks:
+                # A callback consumes per-step metrics: compute the exact
+                # full-batch training loss (separate from the OT cost) and build
+                # the full metrics dict (velocity_mag included).
+                with torch.no_grad():
+                    output = model(inputs)
+                    train_loss = loss_fn(output, targets).item()
+                metrics = {
+                    "step": global_step,
+                    "epoch": epoch,
+                    "loss": train_loss,
+                    "ot_cost": ot_cost,
+                    "displacement": (state.displacement_sqnorms[-1] if state.displacement_sqnorms else 0.0),
+                    "velocity_mag": (torch.norm(state.velocity).item() if state.velocity is not None else 0.0),
+                    "converged": (state.linear_convergence[-1] if state.linear_convergence else False),
+                    "absorb_count": getattr(state, "absorb_count", 0),
+                }
+                for cb in callbacks:
+                    if cb.on_step_end(metrics):
+                        stop = True
+                        break
+            else:
+                # No callback consumes per-step metrics: skip the extra
+                # full-batch forward + its host syncs (.item()/velocity norm)
+                # and track the OT cost as the loss signal for logging / best.
+                train_loss = ot_cost
 
             epoch_loss_sum += train_loss
             epoch_loss_count += 1
 
-            # Build metrics dict
-            state = optimizer.state
-            metrics = {
-                "step": global_step,
-                "epoch": epoch,
-                "loss": train_loss,
-                "ot_cost": state.costs[-1] if state.costs else 0.0,
-                "displacement": (state.displacement_sqnorms[-1] if state.displacement_sqnorms else 0.0),
-                "velocity_mag": (torch.norm(state.velocity).item() if state.velocity is not None else 0.0),
-                "converged": (state.linear_convergence[-1] if state.linear_convergence else False),
-                "absorb_count": getattr(state, "absorb_count", 0),
-            }
-
-            # Invoke on_step_end callbacks
-            for cb in callbacks:
-                if cb.on_step_end(metrics):
-                    stop = True
-                    break
+            # Snapshot the weights at the lowest tracked loss.
+            if config.restore_best and train_loss < best_loss:
+                best_loss = train_loss
+                best_state = copy.deepcopy(model.state_dict())
 
             if stop:
                 break
@@ -218,6 +252,10 @@ def train(
         for cb in callbacks:
             cb.on_epoch_end(epoch_metrics)
 
+    # Restore the best-seen weights (load_state_dict copies in place, keeping dtype/device).
+    if config.restore_best and best_state is not None:
+        model.load_state_dict(best_state)
+
     return model