PyEPO: A PyTorch/JAX-based End-to-End Predict-then-Optimize Tool

PyEPO is a Python-based, open-source software that supports modeling and solving predict-then-optimize problems with linear objective functions. The core capability of PyEPO is to build optimization models with GurobiPy, COPT, Pyomo, Google OR-Tools, MPAX, or any other solvers and algorithms, then embed the optimization model into an artificial neural network for the end-to-end training. For this purpose, PyEPO implements various methods as PyTorch autograd modules, with a mirroring JAX frontend.

Features

Implement SPO+, PG, DPO (additive and multiplicative perturbations), PFYL (additive and multiplicative perturbations), I-MLE, AI-MLE, L2-regularized RFWO/RFYL, DBB, NID, CaVE, NCE, CMAP, and LTR
Support Gurobi, COPT, Pyomo, Google OR-Tools, and MPAX API
Symbolic modeling with pyepo.dsl: define an LP, MIP, or supported fixed-quadratic objective once, then compile it to any backend
JAX frontend (pyepo.func.jax): train any loss in JAX/Flax with jax.grad
Support parallel computing for optimization solvers
Support solution caching to speed up training
Support kNN robust loss to improve decision quality

CaVE for Binary Linear Programs

For end-to-end learning on binary linear programs (TSP, CVRP, knapsack, ...), PyEPO ships CaVE. CaVE replaces the per-step ILP solve with a cone-alignment projection onto the binding-constraint normals at the true optimum, backed by an interior-point QP solver (Clarabel). Because the cone projection is far cheaper than the per-instance ILP solve, CaVE trains an order of magnitude faster than SPO+ at TSP scale.

GPU-Accelerated Solving with MPAX

PyEPO integrates MPAX, a JAX-based mathematical programming solver using the PDHG algorithm for GPU-accelerated optimization. Key advantages: (1) GPU-native solving — the first-order PDHG method runs efficiently on GPU; (2) batch solving — an entire mini-batch can be solved simultaneously via vectorization; (3) no GPU-CPU data transfer overhead — both the neural network and the solver stay on GPU, eliminating the data transfer bottleneck.

Documentation

The official docs can be found at https://khalil-research.github.io/PyEPO.

Publication

PyEPO: A PyTorch-based End-to-End Predict-then-Optimize Library for Linear and Integer Programming (Mathematical Programming Computation)

Metadata

Release files for pyepo 2.2.7

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for pyepo 2.2.7
File	Size	Uploaded
pyepo-2.2.7.tar.gz	167.0 kB	Jul 2, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for pyepo 2.2.7
File	Interpreter	ABI	Platform	Reset
pyepo-2.2.7-py3-none-any.whl 162.7 kB Jul 2, 2026	Python 3	none	any	Details

Total release size: 329.7 kB

Release files / pyepo-2.2.7.tar.gz

Download URL	pyepo-2.2.7.tar.gz
Size	167.0 kB
Tags	Source
SHA-256 checksum How to use checksums	`c4614dcd1e22ffeefa1fd9df5b60f51e6502c423e1fd47f4846804a2c1cea0e2`
BLAKE2b-256 checksum How to use checksums	`ec26d63c43a821d140d3e198ae86c5f0956c9d50abf23252fff77df5ea8b9435`
Upload date	Jul 2, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.13.12`

Release files / pyepo-2.2.7-py3-none-any.whl

Download URL	pyepo-2.2.7-py3-none-any.whl
Size	162.7 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`218db4e6cf4a0c590a9c64683e8d51cc2725796da6d0c1bf1f85e21acee990d3`
BLAKE2b-256 checksum How to use checksums	`b2de66732472e4f75f670791e61988d0861e532e3d504ae80c3a4dd1530e9bb0`
Upload date	Jul 2, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.13.12`