tvm
Loading...
Searching...
No Matches
sorting.h
Go to the documentation of this file.
1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
24#ifndef TVM_RELAX_OP_SORTING_H_
25#define TVM_RELAX_OP_SORTING_H_
26
27#include <tvm/relax/expr.h>
28#include <tvm/tirx/index_map.h>
29
30namespace tvm {
31namespace relax {
32
34struct SortAttrs : public AttrsNode {
35 int axis;
37
38 static void RegisterReflection();
40}; // struct SortAttrs
41
43struct ArgsortAttrs : public AttrsNode {
44 int axis;
46 ffi::Optional<DLDataType> dtype;
47
48 static void RegisterReflection();
50}; // struct ArgsortAttrs
51
53struct TopKAttrs : public AttrsNode {
54 int k;
55 int axis;
56 bool largest;
57 ffi::String ret_type;
58 ffi::Optional<DLDataType> dtype;
59
60 static void RegisterReflection();
62}; // struct TopKAttrs
63
64} // namespace relax
65} // namespace tvm
66
67#endif // TVM_RELAX_OP_SORTING_H_
Base class of all attribute class.
Definition attrs.h:50
Defines a remapping of tensor indices.
An object that builds and maintains block scope and StmtSref mapping for Dependence analysis.
Definition attr.h:23
Attributes used in argsort operator.
Definition sorting.h:43
ffi::Optional< DLDataType > dtype
Definition sorting.h:46
static void RegisterReflection()
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, AttrsNode)
int axis
Definition sorting.h:44
bool descending
Definition sorting.h:45
Attributes used in sort operator.
Definition sorting.h:34
static void RegisterReflection()
bool descending
Definition sorting.h:36
int axis
Definition sorting.h:35
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SortAttrs", SortAttrs, AttrsNode)
Attributes used in topk operator.
Definition sorting.h:53
int k
Definition sorting.h:54
ffi::Optional< DLDataType > dtype
Definition sorting.h:58
ffi::String ret_type
Definition sorting.h:57
int axis
Definition sorting.h:55
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, AttrsNode)
static void RegisterReflection()
bool largest
Definition sorting.h:56