Is there a mechanism to define a custom JIT operator which returns variable number of outputs? (Tuple
of any types)
I tried defining ns::customOp(...) -> (...)
, but this doesn’t seem to work.
When I save a model with a custom operator which returns 2 outputs Tuple[Tensor, Tensor]
and load the saved model, I see error saying the number of outputs mismatch .
Return value was annotated as having type Tuple[Tensor, Tensor] but is actually of type Tuple[]:
My current operator registration is based off of some of the prim
op registrations
RegisterOperators r({Operator(
"ns::customOp(…) -> (…)",
[](const Node* node) -> Operation {
return [node](Stack& stack) {
myLogic(stack, node);
return 0;
};
}, AliasAnalysisKind::FROM_SCHEMA)});
How do we define an operator which is dynamic and can take variable number of inputs and return variable number of outputs?
[Similar to This discussion]