[Bugfix][Op] Register attributes for unique and print (apache#248)

Attempting to use `dump_ast` on functions containing the operators `relax.unique` and `relax.print` previously crashed due to being unable to query their attributes' keys. It turned out that this was a problem with the operator attributes: They had not been registered on the Python side, so Python representation treated them as opaque TVM objects. This PR corrects this mistake.

Author: slyubomirsky

python/tvm/relax/op/op_attrs.py

@@ -32,3 +32,13 @@ class VMAllocStorageAttrs(Attrs):
 @tvm._ffi.register_object("relax.attrs.VMAllocTensorAttrs")
 class VMAllocTensorAttrs(Attrs):
     """Attributes used in VM alloc_tensor operators"""
+
+
+@tvm._ffi.register_object("relax.attrs.UniqueAttrs")
+class UniqueAttrs(Attrs):
+    """Attributes used for the unique operator"""
+
+
+@tvm._ffi.register_object("relax.attrs.PrintAttrs")
+class PrintAttrs(Attrs):
+    """Attributes used for the print operator"""

tests/python/relax/test_ast_printer.py

@@ -336,5 +336,40 @@ def foo(x: Tensor((m, n), "float32")):
     assert call_tir_text in foo_str
 
 
+def test_operators():
+    # the operator attributes need to be registered to work in the printer
+
+    @R.function
+    def foo(x: Tensor):
+        return relax.unique(x, sorted=True)
+
+    foo_str = strip_whitespace(
+        dump_ast(
+            foo,
+            include_type_annotations=False,
+            include_shape_annotations=False,
+        )
+    )
+    # checking that the attributes are present
+    assert '"sorted":1' in foo_str
+    assert '"return_inverse"' in foo_str
+    assert '"return_counts"' in foo_str
+    assert '"dim"' in foo_str
+
+    @R.function
+    def bar(x: Tensor):
+        return relax.print(x, format="{}")
+
+    bar_str = strip_whitespace(
+        dump_ast(
+            bar,
+            include_type_annotations=False,
+            include_shape_annotations=False,
+        )
+    )
+    print_attrs_str = strip_whitespace('{"format": "{}"}')
+    assert print_attrs_str in bar_str
+
+
 if __name__ == "__main__":
     pytest.main([__file__])