24 #ifndef TVM_RELAX_ATTRS_SORTING_H_
25 #define TVM_RELAX_ATTRS_SORTING_H_
40 refl::ObjectDef<SortAttrs>()
42 "Axis along which the sort is computed."
43 "The default the last axis is used.",
44 refl::DefaultValue(-1))
46 "Whether to sort in descending order."
47 "If it is not specified, it defaults to the ascending order.",
48 refl::DefaultValue(
false));
61 refl::ObjectDef<ArgsortAttrs>()
63 "Axis along which the argsort is computed."
64 "The default the last axis is used.",
65 refl::DefaultValue(-1))
67 "Whether to argsort in descending order."
68 "If it is not specified, it defaults to the ascending order.",
69 refl::DefaultValue(
false))
86 refl::ObjectDef<TopKAttrs>()
87 .def_ro(
"k", &
TopKAttrs::k,
"Number of top elements to select")
88 .def_ro(
"axis", &
TopKAttrs::axis,
"Axis along which to sort the input tensor.",
89 refl::DefaultValue(-1))
91 "The return type [both, values, indices]."
92 "both - return both top k data and indices."
93 "values - return top k data only."
94 "indices - return top k indices only.",
95 refl::DefaultValue(
"both"))
97 "Whether to return largest or smallest elements."
98 "By default, return the largest k elements.",
99 refl::DefaultValue(
true))
Adapter for AttrsNode with the new reflection API.
Definition: attrs.h:385
Base class of all attribute class.
Definition: attrs.h:102
Runtime primitive data type.
Definition: data_type.h:47
Defines a remapping of buffer indices.
Definition: repr_printer.h:91
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
DataType NullValue< DataType >()
Definition: attrs.h:61
Attributes used in argsort operator.
Definition: sorting.h:54
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, BaseAttrsNode)
static void RegisterReflection()
Definition: sorting.h:59
DataType dtype
Definition: sorting.h:57
int axis
Definition: sorting.h:55
bool descending
Definition: sorting.h:56
Attributes used in sort operator.
Definition: sorting.h:34
static void RegisterReflection()
Definition: sorting.h:38
bool descending
Definition: sorting.h:36
int axis
Definition: sorting.h:35
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SortAttrs", SortAttrs, BaseAttrsNode)
Attributes used in topk operator.
Definition: sorting.h:77
int k
Definition: sorting.h:78
DataType dtype
Definition: sorting.h:82
ffi::String ret_type
Definition: sorting.h:81
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, BaseAttrsNode)
int axis
Definition: sorting.h:79
static void RegisterReflection()
Definition: sorting.h:84
bool largest
Definition: sorting.h:80