Documentation

Defining a Dynamic Multi-Output Operator (Split)

Some operators allow varied number of outputs. The following takes the Split operator as an example to describe how to define an operator with multiple dynamic outputs.

Split operator prototype definition is as follows.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
REG_OP(Split)
    .INPUT(split_dim, TensorType({DT_INT32}))
    .INPUT(x, TensorType({DT_COMPLEX128, DT_COMPLEX64, DT_DOUBLE, DT_FLOAT,  DT_FLOAT16, DT_INT16,
                          DT_INT32,      DT_INT64,     DT_INT8,   DT_QINT16, DT_QINT32,  DT_QINT8,
                          DT_QUINT16,    DT_QUINT8,    DT_UINT16, DT_UINT32, DT_UINT64,  DT_UINT8,
                          DT_BF16,       DT_BOOL}))
    .DYNAMIC_OUTPUT(y, TensorType({DT_COMPLEX128, DT_COMPLEX64, DT_DOUBLE, DT_FLOAT,  DT_FLOAT16, DT_INT16,
                                   DT_INT32,      DT_INT64,     DT_INT8,   DT_QINT16, DT_QINT32,  DT_QINT8,
                                   DT_QUINT16,    DT_QUINT8,    DT_UINT16, DT_UINT32, DT_UINT64,  DT_UINT8,
                                   DT_BF16,       DT_BOOL}))
    .REQUIRED_ATTR(num_split, Int)
    .OP_END_FACTORY_REG(Split)

The Split operator allows multiple dynamic outputs. Call create_dynamic_output_outputName to create dynamic outputs.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
auto split = op::Split("split")
  .set_input_x(data)                         // For details about how to create a Data operator, see .
  .set_input_split_dim(const)                // For details about how to create a Const operator, see .
  .set_attr_num_split(2)
  .create_dynamic_output_y(2);           // Create dynamic output y of the Split operator, including two outputs.

auto addn = op::AddN("addn")
  .create_dynamic_input_x(1)                 // Create dynamic input x, including one input.
  .set_dynamic_input_x(0, split, "y0")       // Set the first input of the AddN operator. split indicates the input operator, and y0 indicates the first output of the Split operator.
  .set_attr_N(1);                            // Set the value of attribute N to 1, indicating that the operator has one input.

auto softplus = op::Softplus("softplus")
  .set_input_x(split, "y1");          // Set the input of the Softplus operator. split indicates the input operator, and y1 indicates the second output of the Split operator.
Parent topic: Operator Expression