How to use a custom optim shape inference function#

The built-in shape-inference dispatch table covers all standard ONNX operators. Models that include custom operators from a private domain cannot be shape-inferred out of the box: the engine encounters a node it does not recognise and raises ValueError.

onnx-light lets you plug in a callback that handles the shape inference for any (domain, op_type) pair, in Python and in C++. Once registered, the callback is invoked transparently by compute_shape_node() and compute_shape_model() in Python, and by onnx::onnx_optim::shapes::ShapesContext::ComputeShapeNode() and onnx::onnx_optim::shapes::ShapesContext::ComputeShapeModel() in C++.

Import the shape-inference module#

All public symbols are accessible via the high-level Python module, or the matching C++ headers:

from onnx_light.onnx_optim.shape_inference import (
    ShapesContext,
    OptimTensor,
    compute_shape_node,
    compute_shape_model,
)

#include "onnx_optim/optim_tensor.h"
#include "onnx_optim/shapes/shape_inference.h"
#include "onnx_optim/shapes/shapes_context.h"

using namespace onnx::onnx_optim;          // OptimTensor, OptimShape, OptimDim
using namespace onnx::onnx_optim::shapes;  // ShapesContext

Write the callback#

The callback must have the signature fn(ctx, node) -> None in Python and void fn(ShapesContext &ctx, const NodeProto &node) in C++.

  • ctx is the current ShapesContext. Call get() to read input descriptors and set() to write output descriptors.

  • node is the NodeProto being processed. Use node.input, node.output, and node.attribute to inspect the operator’s operands and attributes.

  • The callback must register an OptimTensor for every output of the node before returning.

def infer_my_op(ctx, node):
    """Infers the output shape of MyDomain::MyOp."""
    x = ctx.get(str(node.input[0]))
    # derive the output shape from the input(s) and attributes
    out_shape = list(x.shape)
    ctx.set(str(node.output[0]), OptimTensor(x.dtype, out_shape))

void infer_my_op(ShapesContext &ctx, const NodeProto &node) {
  const OptimTensor &x = ctx.Get(node.input(0).as_string());
  // derive the output shape from the input(s) and attributes
  OptimShape out_shape = x.Shape();
  ctx.Set(node.output(0).as_string(),
          OptimTensor(/*data=*/nullptr, x.Dtype(), out_shape));
}

Symbolic dimensions (string values such as "N" or "batch") are preserved automatically; compare them with isinstance(d, int) before doing arithmetic on them in Python (or OptimDim::IsInt in C++).

Register the callback#

Call set_custom_shape_inference_function() (C++ SetCustomShapeInferenceFunction()) on the context before running inference:

ctx = ShapesContext()
ctx.set_custom_shape_inference_function("my.domain", "MyOp", infer_my_op)

ShapesContext ctx;
ctx.SetCustomShapeInferenceFunction("my.domain", "MyOp", infer_my_op);

Passing an empty string as the domain normalises it to "ai.onnx" (the default ONNX domain). The registration is stored on the context object and applied for every subsequent call to compute_shape_node or compute_shape_model that uses the same context.

You can check whether a callback is registered:

ctx.has_custom_shape_inference_function("my.domain", "MyOp")  # True
list(ctx.custom_shape_inference_keys())  # ["my.domain:MyOp"]

ctx.GetCustomShapeInferenceFunction("my.domain", "MyOp") != nullptr;  // true
for (const auto &kv : ctx.CustomShapeInferenceFunctions()) {
  // kv.first == "my.domain:MyOp"
}

Infer shapes for a single node#

compute_shape_node() (C++ ComputeShapeNode()) processes one node at a time. Seed the context with the node’s input descriptors first:

import onnx_light.onnx as onnxl
import onnx_light.onnx.defs as defs
import onnx_light.onnx.helper as oh

defs.register_onnx_operator_set_schema()

node = oh.make_node("MyOp", ["X"], ["Y"], domain="my.domain")

ctx = ShapesContext()
ctx.set_opset_version("my.domain", 1)
ctx.set("X", OptimTensor(onnxl.TensorProto.FLOAT, [4, 8]))
ctx.set_custom_shape_inference_function("my.domain", "MyOp", infer_my_op)

compute_shape_node(ctx, node)
print(ctx.get("Y").shape)   # [4, 8]

#include "onnx_proto/onnx_helper.h"  // make_node helpers, etc.

NodeProto node;
node.set_op_type("MyOp");
node.set_domain("my.domain");
node.add_input("X");
node.add_output("Y");

ShapesContext ctx;
ctx.SetOpsetVersion("my.domain", 1);
ctx.Set("X", OptimTensor(/*data=*/nullptr, TensorType::kFloat,
                         {OptimDim(4), OptimDim(8)}));
ctx.SetCustomShapeInferenceFunction("my.domain", "MyOp", infer_my_op);

ctx.ComputeShapeNode(node);
// ctx.Get("Y").Shape() == {4, 8}

Infer shapes for a whole model#

compute_shape_model() (C++ ComputeShapeModel()) walks every node of the main graph in topological order. Register the callback on the context before calling it:

ctx = ShapesContext()
ctx.set_custom_shape_inference_function("my.domain", "MyOp", infer_my_op)
compute_shape_model(ctx, model)

for name in ["intermediate", "output"]:
    t = ctx.get(name)
    print(name, list(t.shape), t.dtype)

ShapesContext ctx;
ctx.SetCustomShapeInferenceFunction("my.domain", "MyOp", infer_my_op);
ctx.ComputeShapeModel(model);

for (const char *name : {"intermediate", "output"}) {
  const OptimTensor &t = ctx.Get(name);
  // t.Shape(), t.Dtype()
}

infer_shapes_model() (C++ onnx::onnx_optim::shapes::InferShapesModel()) creates its own internal context and does not accept a pre-configured context. Use compute_shape_model() (C++ ComputeShapeModel()) instead when your model contains custom operators, then apply the inferred shapes back manually if needed:

from onnx_light.onnx_optim.shape_inference import (
    compute_shape_model,
    apply_inferred_shapes_to_model,
)

ctx = ShapesContext()
ctx.set_custom_shape_inference_function("my.domain", "MyOp", infer_my_op)
compute_shape_model(ctx, model)
apply_inferred_shapes_to_model(ctx, model)

ShapesContext ctx;
ctx.SetCustomShapeInferenceFunction("my.domain", "MyOp", infer_my_op);
ctx.ComputeShapeModel(model);
ctx.ApplyInferredShapesToModel(model);

See also#