adding ArgReduceAttr to to inherit from Attrs (apache#9606)

* adding ArgReduceAttrs as python class with register_object wrapper * cleaning
ylc · Jan 7, 2022 · 55fd5e7 · 55fd5e7
1 parent dc03d13
commit 55fd5e7
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/python/tvm/relay/op/op_attrs.py b/python/tvm/relay/op/op_attrs.py
@@ -509,6 +509,11 @@ class ReduceAttrs(Attrs):
     """Attributes used in reduction operators (e.g. sum)"""
 
 
+@tvm._ffi.register_object("relay.attrs.ArgReduceAttrs")
+class ArgReduceAttrs(Attrs):
+    """Attributes used in reduction operators (e.g. argmin/argmax)"""
+
+
 @tvm._ffi.register_object("relay.attrs.VarianceAttrs")
 class VarianceAttrs(Attrs):
     """Attributes used in reduction operators (e.g. sum)"""