Temporary fix of ConicResult.value_of(v) for v that have shapes that are expression

1 files changed, 11 insertions, 2 deletions

diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
index 02d5a1e..2800423 100644
--- a/sumofsquares/problems.py
+++ b/sumofsquares/problems.py
@@ -15,6 +15,7 @@ from typing import Any, Sequence
 from typing_extensions import override
 
 from polymatrix.expression.expression import Expression, VariableExpression
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 from polymatrix.expressionstate import ExpressionState
 from polymatrix.polymatrix.mixins import PolyMatrixMixin
 from polymatrix.polymatrix.index import MonomialIndex
@@ -121,8 +122,16 @@ class ConicResult(Result):
             var = var.underlying
 
         if var not in self.values:
-            raise KeyError(f"There is no result for the variable {var}, "
-                           "was the problem successfully solved?")
+            # FIXME: this is a temporary fix here.
+            if isinstance(var.shape, ExpressionBaseMixin):
+                state = poly.make_state()
+                state, shapepm = var.shape.apply(state)
+                shape = (shapepm.at(0,0).constant(), shapepm.at(1,0).constant())
+                var = replace(var, shape=shape)
+
+                if var not in self.values:
+                    raise KeyError(f"There is no result for the variable {var}. "
+                                   f"Was the problem successfully solved?")
 
         return self.values[var]