diff options
| -rw-r--r-- | polymatrix/expression/to.py | 7 |
1 files changed, 4 insertions, 3 deletions
diff --git a/polymatrix/expression/to.py b/polymatrix/expression/to.py index 2a43c0e..2707fb4 100644 --- a/polymatrix/expression/to.py +++ b/polymatrix/expression/to.py @@ -49,7 +49,8 @@ def to_sympy( expr: Expression, ) -> StateMonadMixin[ExpressionState, sympy.Expr | sympy.Matrix]: - def polymatrix_to_sympy(state: ExpressionState) -> sympy.Expr | sympy.Matrix: + def polymatrix_to_sympy(state: ExpressionState) -> tuple[ExpressionState, sympy.Expr | sympy.Matrix]: + # Convert to polymatrix state, pm = expr.apply(state) @@ -72,7 +73,7 @@ def to_sympy( if math.prod(pm.shape) == 1: # just return the expression - return m[0, 0] + return state, m[0, 0] - return m + return state, m return init_state_monad(polymatrix_to_sympy) |