Custom operator with multiple input/output types

[netaz]

I have an onnxruntime custom operator that can support multiple input and output types. For example, inputs can be FP32, FP16 and BF16.
I can't find a way to express this using ONNXTensorElementDataType GetOutputType which can only return one type.
I could not find examples of documentation on those.

[adrianlizarraga]

Hi @netaz,

If your custom op has 1 input and 1 output that can be of various tensor types, then you can use ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED to indicate that the input (or output) can potentially be of any type defined in the ONNTensorElementDataType enum

Example:

struct MyCustomOp: Ort::CustomOpBase<MyCustomOp, MyCustomKernel> {
  // ...

  size_t GetInputTypeCount() const { return 1; };  // One input only
  ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;  // Can be f32, f16, uint8, etc.
  };

  size_t GetOutputTypeCount() const { return 1; };  // One output only
  ONNXTensorElementDataType GetOutputType(size_t /*index*/) const {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;  // Can be f32, f16, uint8, etc.
  };
}

Alternatively, if your custom op has multiple inputs/outputs with fixed types, then you can return the type based on the input/output index.

Example:

struct MyCustomOp: Ort::CustomOpBase<MyCustomOp, MyCustomKernel> {
  // ...

  size_t GetInputTypeCount() const { return 2; };  // Two inputs
  ONNXTensorElementDataType GetInputType(size_t index) const {
    if (index == 0) {
      return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;  // The first input is always an f32
    }

    return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;  // The second input is always an f16.
  };

  size_t GetOutputTypeCount() const { return 3; };  // Three outputs
  ONNXTensorElementDataType GetOutputType(size_t index) const {
    // Similar code for outputs.
  };
}

Otherwise, if your custom op can have a dynamic number of inputs and outputs of different types (i.e., variadic like printf), then you can define a custom op with variadic inputs and outputs.

Example:

struct MyCustomOp: Ort::CustomOpBase<MyCustomOp, MyCustomKernel> {
  // ...

  size_t GetInputTypeCount() const { return 1; };  // One input
   OrtCustomOpInputOutputCharacteristic GetInputCharacteristic(size_t index) const noexcept {
    return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC;  // The input is variadic (like printf)
  }
  ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;  // Input can be any type
  };
  bool GetVariadicInputHomogeneity() const {
    return false;  // Each input argument can be of different types (heterogeneous).
  }

  size_t GetOutputTypeCount() const { return 1; };  // One output only
   OrtCustomOpInputOutputCharacteristic GetOutputCharacteristic(size_t index) const noexcept {
    return OrtCustomOpInputOutputCharacteristic::INPUT_OUTPUT_VARIADIC;  // The output is variadic
  }
  ONNXTensorElementDataType GetOutputType(size_t /*index*/) const {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;  // Output can be any type
  };
  bool GetVariadicOuputHomogeneity() const {
    return false;  // Each output operand can be of different type (heterogeneous).
  }
}

Please refer to the C API header for more information on variadic inputs/outputs.

[netaz]

Thanks a lot for the information!

[abhishek27m1992github]

@adrianlizarraga still getting error, when i used like mentioned above
size_t GetInputTypeCount() const { return 1; }
ONNXTensorElementDataType GetInputType(size_t /index/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; }; // ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT

    size_t GetOutputTypeCount() const { return 1; }
    ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; } // ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT

getting below error -

Error: /onnxruntime_src/onnxruntime/core/session/custom_ops.cc:890 onnx::OpSchema onnxruntime::CreateSchema(const std::string&, const std::vector<const OrtCustomOp*>&) 1 == undefined was false. There must be one (and only one) dynamic typed input to the custom op. Its type info at runtime will be used to infer the type info of this dynamic typed output which is required for the success of the model loading step. More than one dynamic typed inputs are currently not supported as differing types at runtime means the output type cannot be inferred without which model loading cannot proceed.

[minrui-hust]

if we set output type to ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED, when and how we actualy specify the output type according to input type? In the compute function, we can use GetOutput to get output, but this function do not specify output type, so what is the output value type?

[adrianlizarraga]

Hi @minrui-hust,

The type of an undefined output is inferred from a corresponding input of undefined type. For this reason, a custom operator with an output of undefined type must also have 1 (and only 1) input of undefined type.

Ex:

• A custom op with input tensor(float32) and an "undefined" output type is invalid. ORT requires an input of undefined type from which to infer the output's type
• A custom op with inputs (undefined, tensor(float32)) and an undefined output type is valid. The output's actual type at runtime must match the undefined input's type (which is determined by the upstream node in the graph).

[minrui-hust]

Thanks for your reply.
So the output type inference is done by ORT, user can do nothing about it, unlike tensorRT, shape and type inference is controlled by user.