Allow passing state object to sos.solve_problem

1 files changed, 2 insertions, 4 deletions

diff --git a/sumofsquares/__init__.py b/sumofsquares/__init__.py
index a29b055..e6af1e7 100644
--- a/sumofsquares/__init__.py
+++ b/sumofsquares/__init__.py
@@ -76,9 +76,7 @@ from .abc import Problem, Result, Set, Constraint, Solver
 from .canon import Canonicalization, PutinarPSatz
 from .constraints import BasicSemialgebraicSet, NonNegative
 from .problems import SOSProblem
-from .utils import partition
 from .variable import (
-        OptVariable,
         from_name as internal_from_name,
         from_names as internal_from_names)
 
@@ -125,7 +123,7 @@ def maximize(cost, *args, **kwargs) -> SOSProblem:
     return make_problem(-cost, *args, **kwargs)
 
 
-def solve_problem(*args, verbose: bool = True, **kwargs) -> tuple[Problem, Result]:
+def solve_problem(*args, verbose: bool = True, state: ExpressionState | None = None, **kwargs) -> tuple[Problem, Result]:
     """
     Solve a sum-of-squares optimization problem.
     This function is just a shorthand for
@@ -135,4 +133,4 @@ def solve_problem(*args, verbose: bool = True, **kwargs) -> tuple[Problem, Resul
         result = prob.solve(verbose)
     """
     prob = make_problem(*args, **kwargs)
-    return prob, prob.solve(verbose=verbose)
+    return prob, prob.solve(verbose=verbose, state=state)