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

* adding ArgReduceAttrs as python class with register_object wrapper * cleaning
yangulei · Jan 11, 2022 · e3470d6 · e3470d6
1 parent 3ebabdd
commit e3470d6
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)"""