pyepo 2.2.1__tar.gz → 2.2.2__tar.gz

{pyepo-2.2.1 → pyepo-2.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.1
+Version: 2.2.2
 Summary: PyTorch-based End-to-End Predict-then-Optimize Tool
 Author-email: Bo Tang <bolucas.tang@mail.utoronto.ca>
 License: MIT

{pyepo-2.2.1 → pyepo-2.2.2}/pyepo/metric/regret.py

@@ -19,6 +19,8 @@ from pyepo.model.mpax import optMpaxModel
 from pyepo.utils import _EPS, costToNumpy, getArgs
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from torch import nn
     from torch.utils.data import DataLoader
 
@@ -75,7 +77,7 @@ def _checkLinearObj(optmodel) -> None:
 
 
 def regret(
-    predmodel: nn.Module,
+    predmodel: nn.Module | Callable,
     optmodel: optModel,
     dataloader: DataLoader,
     processes: int = 1,
@@ -91,11 +93,17 @@ def regret(
     z^*(\\mathbf{c}_i)`. With the default ``reduction="normalized"`` the
     result is :math:`\\sum_i l_i / \\sum_i |z^*(\\mathbf{c}_i)|`,
     dimensionless and comparable across problem scales; instances with
-    near-zero true optima inflate the ratio. The predictor is evaluated
-    under ``eval()``; its original mode is restored afterwards.
+    near-zero true optima inflate the ratio. PyTorch predictors are
+    evaluated under ``eval()``; the original mode is restored afterwards.
+
+    ``predmodel`` may also be a plain callable ``f(x: np.ndarray) ->
+    array-like`` for JAX/Flax models; pass a ``functools.partial`` that
+    closes over the current parameter pytree, e.g.
+    ``functools.partial(model.apply, params)``.
 
     Args:
-        predmodel: a regression neural network for cost prediction
+        predmodel: a PyTorch ``nn.Module`` for cost prediction, or a
+            JAX callable ``f(x_numpy) -> cost_array``
         optmodel: a PyEPO optimization model
         dataloader: PyTorch DataLoader over an ``optDataset`` (yielding
             ``(x, c, w, z)`` tuples)
@@ -122,9 +130,12 @@ def regret(
     processes = mp.cpu_count() if processes == 0 else processes
     losses = []
     optsum = 0.0
-    # get device (cpu fallback for parameterless predictors)
-    param = next(predmodel.parameters(), None)
-    device = param.device if param is not None else torch.device("cpu")
+    _is_torch = isinstance(predmodel, torch.nn.Module)
+    # get device (cpu fallback for parameterless predictors; JAX callables always use cpu)
+    device = torch.device("cpu")
+    if _is_torch:
+        param = next(predmodel.parameters(), None)
+        device = param.device if param is not None else torch.device("cpu")
     # multi-core: each worker builds its own optmodel once via the initializer
     pool = None
     if processes > 1:
@@ -134,16 +145,20 @@ def regret(
             initargs=(type(optmodel), getArgs(optmodel)),
         )
     # evaluate under eval(); the original mode is restored afterwards
-    was_training = predmodel.training
-    predmodel.eval()
+    was_training = predmodel.training if _is_torch else False
+    if _is_torch:
+        predmodel.eval()
     try:
         # load data
         for data in dataloader:
             x, c, _, z = data
-            x, c, z = x.to(device), c.to(device), z.to(device)
-            # predict and batch-solve all instances in one call
-            with torch.no_grad():
-                cp = predmodel(x)
+            if _is_torch:
+                x, c, z = x.to(device), c.to(device), z.to(device)
+                with torch.no_grad():
+                    cp = predmodel(x)
+            else:
+                # JAX callable: f(x_numpy) -> array-like
+                cp = torch.as_tensor(np.asarray(predmodel(x.numpy())), dtype=torch.float32)
             # full cost so the MPAX backend can batch-set the objective in one call
             sols, _ = _solve_batch(optmodel._fullCost(cp), optmodel, processes=processes, pool=pool)
             # vectorized regret accumulation (one host sync per batch)
@@ -158,7 +173,8 @@ def regret(
         if pool is not None:
             _close_pool(pool)
         # restore the original mode even if evaluation raises
-        predmodel.train(was_training)
+        if _is_torch:
+            predmodel.train(was_training)
     loss = np.concatenate(losses) if losses else np.empty(0)
     # reduce
     if reduction == "normalized":

{pyepo-2.2.1 → pyepo-2.2.2}/pyepo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.1
+Version: 2.2.2
 Summary: PyTorch-based End-to-End Predict-then-Optimize Tool
 Author-email: Bo Tang <bolucas.tang@mail.utoronto.ca>
 License: MIT

{pyepo-2.2.1 → pyepo-2.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pyepo"
-version = "2.2.1"
+version = "2.2.2"
 description = "PyTorch-based End-to-End Predict-then-Optimize Tool"
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { text = "MIT" }

{pyepo-2.2.1 → pyepo-2.2.2}/tests/test_60_metric.py

@@ -20,7 +20,12 @@ from pyepo.metric.mse import MSE
 from pyepo.metric.regret import _regretFromObj, calRegret
 from pyepo.metric.unambregret import calUnambRegret
 
-from .conftest import NUM_FEAT, LinearPred, requires_gurobi
+from .conftest import (
+    NUM_FEAT, GRID, NUM_DATA, BATCH,
+    LinearPred,
+    requires_gurobi, requires_jax, requires_mpax,
+    _HAS_GUROBI, _HAS_JAX,
+)
 
 
 class _IdentityModel(nn.Module):
@@ -471,6 +476,75 @@ class TestSkScorer:
         assert not np.isclose(expected, swapped)
 
 
+# ============================================================
+# JAX callable path for pyepo.metric.regret
+# ============================================================
+
+@requires_mpax
+class TestDataloaderMetricsJax:
+    """pyepo.metric.regret accepts a plain callable f(x_numpy) -> array."""
+
+    @pytest.fixture(scope="class")
+    def mpax_data(self):
+        import pyepo
+        from pyepo.data.dataset import optDataset
+        from pyepo.model.mpax.shortestpath import shortestPathModel
+
+        x, c = pyepo.data.shortestpath.genData(NUM_DATA, NUM_FEAT, GRID, seed=42)
+        optmodel = shortestPathModel(grid=GRID)
+        dataset = optDataset(optmodel, x, c)
+        loader = DataLoader(dataset, batch_size=BATCH, shuffle=False)
+        return optmodel, dataset, loader
+
+    def test_callable_returns_non_negative_float(self, mpax_data):
+        import pyepo
+        optmodel, _ds, loader = mpax_data
+        reg = pyepo.metric.regret(lambda x: x, optmodel, loader)
+        assert isinstance(reg, float) and reg >= 0
+
+    def test_callable_reductions_consistent(self, mpax_data):
+        import pyepo
+        optmodel, ds, loader = mpax_data
+        fn = lambda x: x  # noqa: E731
+        per = pyepo.metric.regret(fn, optmodel, loader, reduction="none")
+        total = pyepo.metric.regret(fn, optmodel, loader, reduction="sum")
+        assert isinstance(per, np.ndarray) and len(per) == len(ds)
+        assert total == pytest.approx(per.sum(), rel=1e-5)
+
+
+@pytest.mark.skipif(
+    not (_HAS_GUROBI and _HAS_JAX),
+    reason="Parity: Gurobi (exact solve) + JAX both required",
+)
+class TestDataloaderMetricsJaxParity:
+    """JAX callable and torch nn.Module with identical weights give the same regret."""
+
+    def test_same_weights_same_regret(self, sp_data):
+        import functools
+
+        import jax.numpy as jnp
+        import pyepo
+        from flax import linen as nn
+
+        optmodel, _ds, loader = sp_data
+
+        # random torch predictor; capture weights as numpy
+        torch_pred = LinearPred(NUM_FEAT, optmodel.num_cost)
+        torch_pred.eval()
+        w = torch_pred.linear.weight.detach().numpy()   # (num_cost, num_feat)
+        b = torch_pred.linear.bias.detach().numpy()     # (num_cost,)
+
+        # Flax Dense: kernel shape is (num_feat, num_cost) = w.T
+        flax_pred = nn.Dense(optmodel.num_cost)
+        params = {"params": {"kernel": jnp.asarray(w.T), "bias": jnp.asarray(b)}}
+
+        torch_reg = pyepo.metric.regret(torch_pred, optmodel, loader)
+        jax_reg = pyepo.metric.regret(
+            functools.partial(flax_pred.apply, params), optmodel, loader
+        )
+        assert jax_reg == pytest.approx(torch_reg, abs=1e-4)
+
+
 class TestAutoSkScorer:
 
     def test_raises_without_autosklearn(self):