1 files changed, 1 insertions, 7 deletions
diff --git a/polymatrix/expression/mixins/fromstatemonad.py b/polymatrix/expression/mixins/fromstatemonad.py
index a5970b1..254d125 100644
--- a/polymatrix/expression/mixins/fromstatemonad.py
+++ b/polymatrix/expression/mixins/fromstatemonad.py
@@ -27,13 +27,7 @@ class FromStateMonadMixin(ExpressionBaseMixin):
         state, expr = self.monad.apply(state)
         # Case when monad wraps function 
         #  f: ExpressionState -> (State, Expression)
-        if isinstance(expr, Expression):
-            return expr.underlying.apply(state)
-
-        # Case when monad wraps function 
-        #  f: ExpressionState -> (State, M)
-        # where M is a mixin class
-        elif isinstance(expr, ExpressionBaseMixin):
+        if isinstance(expr, Expression | ExpressionBaseMixin):
             return expr.apply(state)
 
         # Case when monad wraps function