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));
51 static constexpr
const char*
_type_key =
"relax.attrs.SortAttrs";
63 refl::ObjectDef<ArgsortAttrs>()
65 "Axis along which the argsort is computed."
66 "The default the last axis is used.",
67 refl::DefaultValue(-1))
69 "Whether to argsort in descending order."
70 "If it is not specified, it defaults to the ascending order.",
71 refl::DefaultValue(
false))
76 static constexpr
const char*
_type_key =
"relax.attrs.ArgsortAttrs";
90 refl::ObjectDef<TopKAttrs>()
91 .def_ro(
"k", &
TopKAttrs::k,
"Number of top elements to select")
92 .def_ro(
"axis", &
TopKAttrs::axis,
"Axis along which to sort the input tensor.",
93 refl::DefaultValue(-1))
95 "The return type [both, values, indices]."
96 "both - return both top k data and indices."
97 "values - return top k data only."
98 "indices - return top k indices only.",
99 refl::DefaultValue(
"both"))
101 "Whether to return largest or smallest elements."
102 "By default, return the largest k elements.",
103 refl::DefaultValue(
true))
108 static constexpr
const char*
_type_key =
"relax.attrs.TopKAttrs";
Adapter for AttrsNode with the new reflection API.
Definition: attrs.h:384
Base class of all attribute class.
Definition: attrs.h:103
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:56
static constexpr const char * _type_key
Definition: sorting.h:76
static void RegisterReflection()
Definition: sorting.h:61
TVM_FFI_DECLARE_FINAL_OBJECT_INFO(ArgsortAttrs, BaseAttrsNode)
DataType dtype
Definition: sorting.h:59
int axis
Definition: sorting.h:57
bool descending
Definition: sorting.h:58
Attributes used in sort operator.
Definition: sorting.h:34
static void RegisterReflection()
Definition: sorting.h:38
bool descending
Definition: sorting.h:36
TVM_FFI_DECLARE_FINAL_OBJECT_INFO(SortAttrs, BaseAttrsNode)
int axis
Definition: sorting.h:35
static constexpr const char * _type_key
Definition: sorting.h:51
Attributes used in topk operator.
Definition: sorting.h:81
int k
Definition: sorting.h:82
String ret_type
Definition: sorting.h:85
DataType dtype
Definition: sorting.h:86
TVM_FFI_DECLARE_FINAL_OBJECT_INFO(TopKAttrs, BaseAttrsNode)
static constexpr const char * _type_key
Definition: sorting.h:108
int axis
Definition: sorting.h:83
static void RegisterReflection()
Definition: sorting.h:88
bool largest
Definition: sorting.h:84