文档

Defining a Basic Compute Operator (SoftmaxV2)

The following takes the SoftmaxV2 operator as an example to describe how to define a basic operator.

SoftmaxV2 operator prototype definition is as follows.

1
2
3
4
5
6
REG_OP(SoftmaxV2)
    .INPUT(x, TensorType({ DT_DOUBLE, DT_FLOAT16, DT_BF16, DT_FLOAT }))
    .OUTPUT(y, TensorType({ DT_DOUBLE, DT_FLOAT16, DT_BF16, DT_FLOAT }))
    .ATTR(axes, ListInt, {-1})
    .ATTR(half_to_float, Bool, false)
    .OP_END_FACTORY_REG(SoftmaxV2)

In this example, the SoftmaxV2 operator has a required input named x. The following shows the creation of a SoftmaxV2 operator instance.

1
2
auto softmax = op::SoftmaxV2("Softmax")     // Create an operator instance and pass the operator name (for example, Softmax) to this call.
    .set_input_x(matmul2);                // Pass matmul2 as the operator input.
Parent topic: Operator Expression