Class ArgMaxFusion
Defined in File arg_max_fusion.h
Inheritance Relationships
Base Type
public mindspore::ops::ArgMax
(Class ArgMax)
Class Documentation
-
class ArgMaxFusion : public mindspore::ops::ArgMax
ArgMaxFusion defined ArgMax operator prototype of lite.
Public Functions
-
inline ArgMaxFusion()
Constructor.
-
void Init(const bool keep_dims, const bool out_max_value, const int64_t top_k, const int64_t axis = -1)
Method to init the op’s attributes.
- Parameters
keep_dims – [in] Define a boolean value to indicate the dimension of output is equal to that of input or not.
out_max_value – [in] Define a boolean value to indicate whether to output maximum value.
top_k – [in] Define the number of maximum value along with axis.
axis – [in] Define where the argmax operation applies to.
-
void set_keep_dims(const bool keep_dims)
Method to set keep_dims attribute.
- Parameters
keep_dims – [in] Define a boolean value to indicate the dimension of output is equal to that of input or not.
-
void set_out_max_value(const bool out_max_value)
Method to set out_max_value attribute.
- Parameters
out_max_value – [in] Define a boolean value to indicate whether to output maximum value.
-
void set_top_k(const int64_t top_k)
Method to set top_k attribute, default is 1.
- Parameters
top_k – [in] Define the number of maximum value along with axis.
-
bool get_keep_dims() const
Method to get keep_dims attribute.
- Returns
a boolean value to indicate the dimension of output is equal to that of input or not.
-
bool get_out_max_value() const
Method to get out_max_value attribute.
- Returns
a boolean value to indicate whether to output maximum value.
-
int64_t get_top_k() const
Method to get top_k attribute.
- Returns
the number of maximum value along with axis.
-
inline ArgMaxFusion()