Expose a Python interface for inference functions

[jbachurski]

Exposes a node-level API for performing type/shape inference to Python bindings.

Description
(Updated)

• Adds a low-level function OpSchema._infer_node_outputs, to which serialised protobuf bytes are passed - arguments: schema, node, input types, input data/other inference context arguments, returns: inferred output types.
• Adds a public interface shape_inference.infer_node_outputs:

def infer_node_outputs(
     schema: onnx.defs.OpSchema,
     node: onnx.NodeProto,
     input_types: Dict[str, onnx.TypeProto],
     input_data: Optional[Dict[str, onnx.TensorProto]] = None,
     input_sparse_data: Optional[Dict[str, onnx.SparseTensorProto]] = None,
 ) -> Dict[str, onnx.TypeProto]:  # output_types

• Underneath this performs validation of passed in nodes/types, ValidationErrors and InferenceErrors are raised.
• Should be the only used function, OpSchema._infer_node_outputs is internal to the implementation.
• input_types keys are named like node declares its inputs.
• input_data and input_sparse_data are used for passing in known constant inputs (as per C++ usage).
• Could implement passing in partial data propagation results and GraphInferenceContext in the future.
• Adds test cases for some inference calls from the Python side.

Motivation and Context

• See Exposing a Python interface for calling inference functions #4405

[jbachurski]

@gramalingam Would you mind viewing the proposal if this sort of interface would work? In the future it could be extended to support the other arguments, they could be optional so it's backward compatible.

[gramalingam]

Updating the pyi definition for OpSchema would be useful.

[jbachurski]

I think it might be useful to promote an API like this as a wrapper around the raw C++ method (that is, make if part of the public API instead of something in just this test-file). But, for that, I don't think we would want node to be optional. And am not sure about num_outputs either.

When it comes to the Python-side interface: A node has input/output names, operator name & domain, attributes. Own name and docstring don't matter. We are missing basically only the operator version it seems, and the input types. I see some variations of the possible signature:

• Inputs & outputs
  • list[TypeProto] - we already have an ordering of fields inside the node. Node inputs/outputs are positional anyway, so making it explicit with ordering in a list seems okay.
  • dict[str, TypeProto] - since the node names the fields in definition. On the other hand since the node has to get created explicitly first anyway in this approach we might as well use the names that are within the node. I think this shouldn't mess with generics etc. since if we use a name multiple times is has to have the same type - we get some very simple consistency check.
• Operator version
  • version: int - get the current operator schema from a list at this current version based on what is in the node.
  • OpSchema- make the user explicitly provide the schema themselves. Is more redundant, since the Node already has the domain and identifier defined.

Also, where would we put this function? shape_inference would be fine, but really it should have been called inference (as it also infers types and other stuff?). Could also put it in helper. Any ideas?

Let me know your thoughts and I'll get started on some improvements in the Python-interface area.

[jbachurski]

@gramalingam I pushed some changes:

std::unordered_map<std::string, py::bytes> CallNodeInferenceFunction(
    OpSchema* schema,
    const py::bytes& nodeBytes,
    std::unordered_map<std::string, py::bytes> valueTypesByNameBytes,
    std::unordered_map<std::string, py::bytes> inputDataByNameBytes,
    std::unordered_map<std::string, py::bytes> inputSparseDataByNameBytes) {

• Now returns a map, since the inputs are a map already. Doesn't return entries were the output TypeProto was not initialised (inference failed implicitly?).
• Added OpSchema::Verify before the call to the inference (checks if the count of inputs/outputs is correct, etc.), and OpSchema::CheckInputOutputType afterwards (checks if the types are consistent with the signature). Both should throw ValidationError.
• At the same time, inference itself may throw InferenceError. Interestingly, this may also throw errors related to attributes (so things like Cast will check here if exactly one of the right attributes is set).
• Let me know if I should handle exceptions there in some more involved way. I think what is thrown already reflects most of what check_model and infer_shapes do.

# class OpSchema:
    def infer_node_outputs(self, node_proto: bytes, value_types: Dict[str, bytes],
                           input_data: Dict[str, bytes], input_sparse_data: Dict[str, bytes]
                           ) -> Dict[str, bytes]: ...

• This is just a binding for the C++ function. I added it to defs.pyi as above.
• My general feeling is that the attributes after value_types should have empty defaults so that in the future generatedShapeData and graphInferenceContext may be supported as well (but this involves, I think, significant work).

# shape_inference.py
def infer_node_outputs(
    schema: onnx.defs.OpSchema, node: onnx.NodeProto, input_types: Dict[str, onnx.TypeProto],
    input_data: Optional[Dict[str, onnx.TensorProto]] = None,
    input_sparse_data: Optional[Dict[str, onnx.SparseTensorProto]] = None
) -> Dict[str, onnx.TypeProto]: ...

• The public-facing Python function does all of the conversions and expects proper types.
• Already has default arguments, which we may change in the future. Reflect the actual underlying implementation.
• Does some checks if all the expected keys are present and only copies over what is actually required on C++ side for some performance.

Let me know what you think of these changes! I also added some more tests.

[xadupre]

None is better than an empty container as a default value.

[jbachurski]

I don't think I know what you mean here. This is C++ side, so there's no None. Do you mean to make it a pointer and use nullptr/use std::optional?
I think that this wouldn't be a lot better as the map isn't mutated and also in this it makes sense that the default is "no additional information" (for the inference context we have to pass in some map, and it makes sense to make it empty as there's no input data we know).
Could you provide a code example for the change?

[gramalingam]

Looks great, thanks very much!

[jbachurski]

Thanks!
Should be ready for merge, then?

[justinchuby]

Do we explicitly require a dictionary for these params? If not prefer a more generic type like Mapping.

——-

Use more generic types for input parameters and specific types for return values. For example, a function may take a Sequence and return a List.

Diagram for reference: https://gist.github.com/justinchuby/4021cebe9e093f636759a88de325c85f