Repository: microsoft/onnxscript

Issue:
Python bool ops have >=2 operands
`and` and `or` are not binary operators in Python and can have any number of operands. Python coalesces the bool ops into a single AST node when the expression chain semantically allows for it. As a result, we are not generating the correct graph for >2 operands:

```python
from onnxscript import script, opset18, INT32, BOOL

@script(default_opset=opset18)
def is1234(x: INT32) -> BOOL:
  return x == 1 or x == 2 or x == 3 or x == 4
```

Results in only the first two operands for `or` in the graph:

<img width="447" alt="image" src="https://user-images.githubusercontent.com/49539/210890485-02448269-fc3b-4aa9-bca1-6b3497badd82.png">

The bug is in [converter.py](https://github.com/microsoft/onnx-script/blob/a733ec2c435807c21e1edc3bc6212b9c1134cd95/onnxscript/converter.py#L780-L783). All values in the AST node need to be accounted for.

The resulting fix would also support #276.


Discussion:

Resolving patch:
diff --git a/onnxscript/converter.py b/onnxscript/converter.py
--- a/onnxscript/converter.py
+++ b/onnxscript/converter.py
@@ -456,15 +456,17 @@ def translate_docstring(self, node):
             f"Unexpected type {type(node)!r} for node. Unsupoorted version of python."
         )
 
-    def translate_expr(self, node, target="tmp"):
+    def translate_expr(self, node, target="tmp") -> ConverterExpression:
         """Expression-translation generates "IR statements/nodes" that compute the value of
         the expression into a target-variable, and returns the variable that is
         assigned this value.
         """
         if isinstance(node, ast.Call):
             r = self.translate_call_expr(node)
-        elif isinstance(node, (ast.BinOp, ast.BoolOp, ast.BitAnd, ast.BitOr)):
+        elif isinstance(node, (ast.BinOp, ast.BitAnd, ast.BitOr)):
             r = self.translate_bin_op_expr(node)
+        elif isinstance(node, ast.BoolOp):
+            r = self.translate_bool_op_expr(node)
         elif isinstance(node, ast.UnaryOp):
             r = self.translate_unary_op_expr(node)
         elif isinstance(node, ast.Compare):
@@ -759,7 +761,25 @@ def _cast_like_binary_expression(self, op, left, right):
         schema = op.get_schema()
         return autocast.static_cast_inputs(self, schema, left, right)
 
-    def translate_bin_op_expr(self, node):
+    def translate_bool_op_expr(self, node: ast.BoolOp) -> ConverterExpression:
+        if isinstance(node.op, ast.And):
+            op = values.Op(self.default_opset, "And")
+        elif isinstance(node.op, ast.Or):
+            op = values.Op(self.default_opset, "Or")
+        else:
+            raise ValueError(self.message(node, f"Unsupported operator {node.op!r}."))
+
+        expr = self.translate_expr(node.values[0])
+        for operand in node.values[1:]:
+            left, right = self._cast_like_binary_expression(
+                op, expr, self.translate_expr(operand)
+            )
+            ovar = self.generate_unique_name()
+            self.emit([ovar], op, [left, right], [])
+            expr = ConverterExpression(ovar, ConverterExpressionKind.ANY)
+        return expr
+
+    def translate_bin_op_expr(self, node: ast.BinOp):
         op = type(node.op)
         if op not in primop_map:
             raise ValueError(self.message(node, f"Unsupported operator {op!r}."))
@@ -772,18 +792,10 @@ def translate_bin_op_expr(self, node):
             if isinstance(cst, float):
                 attr = [self.ir_builder.make_attr("fmod", 1)]
 
-        opname = primop_map[op]
-        if hasattr(node, "left"):
-            # operation
-            left = self.translate_expr(node.left)
-            right = self.translate_expr(node.right)
-        else:
-            # and, or
-            left = self.translate_expr(node.values[0])
-            right = self.translate_expr(node.values[1])
-
-        op = values.Op(self.default_opset, opname)
-        left, right = self._cast_like_binary_expression(op, left, right)
+        op = values.Op(self.default_opset, primop_map[op])
+        left, right = self._cast_like_binary_expression(
+            op, self.translate_expr(node.left), self.translate_expr(node.right)
+        )
         return op, [left, right], attr
 
     def translate_unary_op_expr(self, node):


Test patch:
diff --git a/onnxscript/test/operator_test.py b/onnxscript/test/operator_test.py
--- a/onnxscript/test/operator_test.py
+++ b/onnxscript/test/operator_test.py
@@ -9,7 +9,7 @@
 
 from onnxscript import script
 from onnxscript.onnx_opset import opset15 as op
-from onnxscript.onnx_types import FLOAT
+from onnxscript.onnx_types import BOOL, FLOAT
 from onnxscript.test.common import testutils
 
 
@@ -42,6 +42,49 @@ def implicit_plus1(A: FLOAT["N"]) -> FLOAT["N"]:  # noqa: F821
 
         self.assertSame(explicit_plus1, implicit_plus1)
 
+    def test_bool_ops_binary(self):
+        @script(default_opset=op)
+        def py_two_or(a: BOOL, b: BOOL) -> BOOL:
+            return a or b
+
+        @script()
+        def onnx_two_or(a: BOOL, b: BOOL) -> BOOL:
+            return op.Or(a, b)
+
+        self.assertSame(py_two_or, onnx_two_or)
+
+        @script(default_opset=op)
+        def py_two_and(a: BOOL, b: BOOL) -> BOOL:
+            return a and b
+
+        @script()
+        def onnx_two_and(a: BOOL, b: BOOL) -> BOOL:
+            return op.And(a, b)
+
+        self.assertSame(py_two_and, onnx_two_and)
+
+    def test_bool_ops_three(self):
+        @script(default_opset=op)
+        def py_three_or(a: BOOL, b: BOOL, c: BOOL) -> BOOL:
+            return a or b or c
+
+        @script()
+        def onnx_three_or(a: BOOL, b: BOOL, c: BOOL) -> BOOL:
+            return op.Or(op.Or(a, b), c)
+
+        self.assertSame(py_three_or, onnx_three_or)
+
+    def test_bool_ops_four(self):
+        @script(default_opset=op)
+        def py_four_and(a: BOOL, b: BOOL, c: BOOL, d: BOOL) -> BOOL:
+            return a and b and c and d
+
+        @script()
+        def onnx_four_and(a: BOOL, b: BOOL, c: BOOL, d: BOOL) -> BOOL:
+            return op.And(op.And(op.And(a, b), c), d)
+
+        self.assertSame(py_four_and, onnx_four_and)
+
 
 if __name__ == "__main__":
     unittest.main()