Class ArgMaxFusion

Inheritance Relationships

Base Type

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.