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

* adding ArgReduceAttrs as python class with register_object wrapper * cleaning
ylc · Jan 13, 2022 · aee4acb · aee4acb
1 parent 3dd4c84
commit aee4acb
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)"""