Repository: microsoft/onnxscript

Issue:
Support bool as an attribute type
#281 

```python
@script()
def aten_logsumexp(self: TReal, dim: INT64["M"], keepdim: bool = False) -> TReal:
    # logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor

    return op.ReduceLogSumExp(self, dim, keepdims=keepdim)
```

produces

ERROR onnxscript/test/function_libs/torch_aten/ops_correctness_test.py - TypeError: value "<onnxscript.values.Dynamic object at 0x7f2611f68640>" is not valid attribute data type.

Even though `dim` should be an input?

cc @gramalingam @fatcat-z @xiaowuhu 


Discussion:
Quick question: is this with the latest opset of ONNX, which updates ReduceLogSumExp?
I was using opset18, onnx 1.13 and onnxruntime 1.13 I think
Ok, the problem is with keepdim. Currently, onnxscript doesn't recognize the "bool" type as a valid attribute type. (ONNX doesn't have bool attributes, it uses int attributes even for booleans.) So, if I use `int` it seems to work for now. We could extend the translator to support bool attributes and map them automatically to int attributes.
I see! That makes a lot of sense. I like the idea of supporting bool since it is a kind of int even in Python 

Resolving patch:
diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py
--- a/onnxscript/function_libs/torch_aten/ops/core.py
+++ b/onnxscript/function_libs/torch_aten/ops/core.py
@@ -11,7 +11,7 @@
 """
 from __future__ import annotations
 
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Optional, Sequence, Tuple, Union
 
 from onnxscript import BOOL, DOUBLE, FLOAT, INT16, INT32, INT64
 from onnxscript.function_libs.torch_aten.registration import torch_op
@@ -4754,12 +4754,21 @@ def aten_to_sparse_csr(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
+@torch_op("aten::topk")
 def aten_topk(
-    self: TensorType, k: int, dim: int = -1, largest: bool = True, sorted: bool = True
-) -> tuple[TensorType, TensorType]:
+    self: TReal, k: INT64, dim: int = -1, largest: bool = True, sorted: bool = True
+) -> Tuple[TReal, INT64]:
     # topk(Tensor self, int k, int dim=-1, bool largest=True, bool sorted=True) -> (Tensor values, Tensor indices)
 
-    raise NotImplementedError()
+    self_is_scalar = op.Size(op.Shape(self)) == 0
+    if self_is_scalar:
+        self = op.Unsqueeze(self, op.Constant(value_ints=[0]))
+    k = op.Reshape(op.Cast(k, to=INT64.dtype), op.Constant(value_ints=[1]))
+    values, indices = op.TopK(self, k, axis=dim, largest=largest, sorted=sorted)
+    if self_is_scalar:
+        values = op.Squeeze(values, op.Constant(value_ints=[0]))
+        indices = op.Squeeze(indices, op.Constant(value_ints=[0]))
+    return values, indices
 
 
 def aten_trace(self: TensorType) -> TensorType:
diff --git a/onnxscript/type_annotation.py b/onnxscript/type_annotation.py
--- a/onnxscript/type_annotation.py
+++ b/onnxscript/type_annotation.py
@@ -17,12 +17,14 @@
     float: onnx.AttributeProto.FLOAT,
     int: onnx.AttributeProto.INT,
     str: onnx.AttributeProto.STRING,
+    bool: onnx.AttributeProto.INT,
 }
 
 _LISTTYPE_TO_ATTRTYPE_MAP = {
     float: onnx.AttributeProto.FLOATS,
     int: onnx.AttributeProto.INTS,
     str: onnx.AttributeProto.STRINGS,
+    bool: onnx.AttributeProto.INTS,
 }
 
 _LIST_CONSTRUCTORS = frozenset([list, typing.List, typing.Sequence, collections.abc.Sequence])


Test patch:
diff --git a/onnxscript/test/common/testutils.py b/onnxscript/test/common/testutils.py
--- a/onnxscript/test/common/testutils.py
+++ b/onnxscript/test/common/testutils.py
@@ -16,6 +16,8 @@ def function_proto(f):
         return f
     if isinstance(f, OnnxFunction):
         return f.to_function_proto()
+    if isinstance(f, str):
+        return parser.parse_function(f)
     raise TypeError(f"Cannot convert {type(f)} to FunctionProto")
 
 
@@ -51,3 +53,6 @@ def assertSame(self, fn1, fn2):
 
     def assertSameGraph(self, graph1, graph2):
         self.assertTrue(isomorphic(graph_proto(graph1), graph_proto(graph2)))
+
+    def assertSameFunction(self, fn1, fn2):
+        self.assertTrue(isomorphic(function_proto(fn1), function_proto(fn2)))
diff --git a/onnxscript/test/function_libs/torch_aten/ops_correctness_test.py b/onnxscript/test/function_libs/torch_aten/ops_correctness_test.py
--- a/onnxscript/test/function_libs/torch_aten/ops_correctness_test.py
+++ b/onnxscript/test/function_libs/torch_aten/ops_correctness_test.py
@@ -12,6 +12,7 @@
 import torch
 from torch.testing._internal import common_device_type, common_methods_invocations
 from torch.testing._internal.opinfo import core as opinfo_core
+from torch.utils import _pytree as pytree
 
 import onnxscript
 from onnxscript.function_libs.torch_aten.ops import core as core_ops
@@ -197,6 +198,20 @@ def _logcumsumexp_input_wrangler(
     return args, kwargs
 
 
+def _topk_input_wrangler(
+    args: list[Any], kwargs: dict[str, Any]
+) -> tuple[list[Any], dict[str, Any]]:
+    # TODO(#305): Sole purpose is to workaround attributes must be in kwargs in onnxscript.
+
+    if len(args) >= 3:
+        kwargs["dim"] = args.pop(2)
+    if len(args) >= 3:
+        kwargs["largest"] = args.pop(2)
+    if len(args) >= 3:
+        kwargs["sorted"] = args.pop(2)
+    return args, kwargs
+
+
 # Ops to be tested for numerical consistency between onnx and pytorch
 # Find the names of the OpInfos in torch/testing/_internal/common_methods_invocations.py
 OPINFO_FUNCTION_MAPPING: dict[
@@ -294,6 +309,10 @@ def _logcumsumexp_input_wrangler(
     "tan": core_ops.aten_tan,
     "tanh": core_ops.aten_tanh,
     "transpose": core_ops.aten_transpose,
+    "topk": (
+        core_ops.aten_topk,
+        _topk_input_wrangler,
+    ),
     "unsqueeze": core_ops.aten_unsqueeze,
     "view": core_ops.aten_view,
     "where": core_ops.aten_where,
@@ -543,28 +562,38 @@ def test_output_match(self, device: str, dtype: torch.dtype, op):
                 torch_output = op(*inputs, **cpu_sample.kwargs)
                 function_output = onnx_function(*input_onnx, **kwargs_onnx)
 
-                if dtype == torch.float32:
-                    # Relax atol and rtol for float32 based on empirical results
-                    # The current most relaxed values are for aten::matmul
-                    rtol = 3.7e-6
-                    atol = 1.8e-5
-                else:
-                    rtol = None
-                    atol = None
-
-                if not isinstance(function_output, np.ndarray):
-                    # An onnxscript tensor
-                    function_output = function_output.value
-
-                # Use torch.testing as opposed to np.testing to ensure dtypes and shapes match
-                torch.testing.assert_close(
-                    torch.tensor(function_output),
-                    torch_output
-                    if isinstance(torch_output, torch.Tensor)
-                    else torch.tensor(torch_output),
-                    rtol=rtol,
-                    atol=atol,
-                )
+                # TODO: add pytree structure comparison.
+                flattened_torch_outputs, _ = pytree.tree_flatten(torch_output)
+                flattened_function_outputs, _ = pytree.tree_flatten(function_output)
+
+                assert len(flattened_torch_outputs) > 0
+                assert len(flattened_torch_outputs) == len(flattened_function_outputs)
+
+                for torch_output, function_output in zip(
+                    flattened_torch_outputs, flattened_function_outputs
+                ):
+                    if dtype == torch.float32:
+                        # Relax atol and rtol for float32 based on empirical results
+                        # The current most relaxed values are for aten::matmul
+                        rtol = 3.7e-6
+                        atol = 1.8e-5
+                    else:
+                        rtol = None
+                        atol = None
+
+                    if not isinstance(function_output, np.ndarray):
+                        # An onnxscript tensor
+                        function_output = function_output.value
+
+                    # Use torch.testing as opposed to np.testing to ensure dtypes and shapes match
+                    torch.testing.assert_close(
+                        torch.tensor(function_output),
+                        torch_output
+                        if isinstance(torch_output, torch.Tensor)
+                        else torch.tensor(torch_output),
+                        rtol=rtol,
+                        atol=atol,
+                    )
 
 
 common_device_type.instantiate_device_type_tests(
diff --git a/onnxscript/test/type_annotation_test.py b/onnxscript/test/type_annotation_test.py
--- a/onnxscript/test/type_annotation_test.py
+++ b/onnxscript/test/type_annotation_test.py
@@ -5,6 +5,9 @@
 
 import unittest
 
+import onnx
+from packaging.version import Version
+
 from onnxscript import script
 from onnxscript.onnx_opset import opset15 as op
 from onnxscript.onnx_types import FLOAT
@@ -66,6 +69,28 @@ def unknown_rank(A: FLOAT[...], B: FLOAT[...]) -> FLOAT[...]:
         with self.assertRaises(ValueError):
             FLOAT[10][20]  # Invalid usage. pylint: disable=pointless-statement
 
+    @unittest.skipIf(
+        Version(onnx.__version__) < Version("1.13"),
+        reason="module 'onnx.parser' has no attribute 'parse_function'",
+    )
+    def test_type_annotation_with_bool_type_for_attribute(self):
+        @script()
+        def bool_type_for_attribute(self: FLOAT[...], sorted: bool) -> FLOAT[...]:
+            out = op.Unique(self, sorted=sorted)
+            return out
+
+        bool_type_for_attribute_txt = """
+            <
+                domain: "this",
+                opset_import: ["": 15]
+            >
+            bool_type_for_attribute <sorted>(self) => (out) {
+                out = Unique <sorted: int = @sorted> (self)
+            }
+
+        """
+        self.assertSameFunction(bool_type_for_attribute, bool_type_for_attribute_txt)
+
 
 if __name__ == "__main__":
     unittest.main()