A model with two function nodes fails to pass ONNXchecker after ONNX infer_shape

[titaiwangms]

A model with two function nodes fails to pass ONNXchecker after ONNX infer_shape.

Error messages:

C.check_model(protobuf_string, full_check)
onnx.onnx_cpp2py_export.checker.ValidationError: Field 'shape' of 'type' is required but missing.

ONNX model proto:

ir_version: 7
producer_name: "pytorch"
producer_version: "2.0.0"
graph {
  node {
    output: "1"
    name: "Constant_0"
    op_type: "Constant"
    attribute {
      name: "value"
      t {
        data_type: 7
        raw_data: "\001\000\000\000\000\000\000\000"
      }
      type: TENSOR
    }
    doc_string: ""
  }
  node {
    input: "arg0_1"
    input: "1"
    output: "3"
    name: "aten_add_1"
    op_type: "aten_add"
    doc_string: ""
    domain: "onnx-script"
  }
  node {
    input: "3"
    output: "4"
    name: "aten_sigmoid_2"
    op_type: "aten_sigmoid"
    doc_string: ""
    domain: "onnx-script"
  }
  name: "torch_jit"
  input {
    name: "arg0_1"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 1
          }
          dim {
            dim_value: 4
          }
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 3
          }
        }
      }
    }
  }
  output {
    name: "4"
    type {
      tensor_type {
        elem_type: 1
      }
    }
  }
  value_info {
    name: "3"
    type {
      tensor_type {
        elem_type: 1
      }
    }
  }
  value_info {
    name: "1"
    type {
      tensor_type {
        elem_type: 7
        shape {
        }
      }
    }
  }
}
opset_import {
  version: 14
}
opset_import {
  domain: "onnx-script"
  version: 1
}
functions {
  name: "aten_add"
  input: "self"
  input: "other"
  output: "return_val"
  attribute: "alpha"
  node {
    output: "alpha"
    name: "n0"
    op_type: "Constant"
    attribute {
      name: "value_float"
      type: FLOAT
      ref_attr_name: "alpha"
    }
    domain: ""
  }
  node {
    output: "tmp"
    name: "n1"
    op_type: "Constant"
    attribute {
      name: "value"
      t {
        data_type: 7
        int64_data: 1
        name: "tmp"
      }
      type: TENSOR
    }
    domain: ""
  }
  node {
    input: "tmp"
    input: "alpha"
    output: "tmp_cast"
    name: "n2"
    op_type: "CastLike"
    domain: ""
  }
  node {
    input: "alpha"
    input: "tmp_cast"
    output: "tmp_0"
    name: "n3"
    op_type: "Equal"
    domain: ""
  }
  node {
    input: "tmp_0"
    output: "cond"
    name: "n4"
    op_type: "Not"
    domain: ""
  }
  node {
    input: "cond"
    output: "other_4"
    name: "n5"
    op_type: "If"
    attribute {
      name: "then_branch"
      g {
        node {
          output: "alpha_1"
          name: "n0"
          op_type: "Constant"
          attribute {
            name: "value_float"
            type: FLOAT
            ref_attr_name: "alpha"
          }
          domain: ""
        }
        node {
          input: "other"
          input: "alpha_1"
          output: "other_2"
          name: "n1"
          op_type: "Mul"
          domain: ""
        }
        name: "thenGraph_4"
        output {
          name: "other_2"
          type {
          }
        }
      }
      type: GRAPH
    }
    attribute {
      name: "else_branch"
      g {
        node {
          input: "other"
          output: "other_3"
          name: "n0"
          op_type: "Identity"
          domain: ""
        }
        name: "elseGraph_4"
        output {
          name: "other_3"
          type {
          }
        }
      }
      type: GRAPH
    }
    domain: ""
  }
  node {
    input: "self"
    input: "other_4"
    output: "return_val"
    name: "n6"
    op_type: "Add"
    domain: ""
  }
  opset_import {
    domain: ""
    version: 18
  }
  domain: "onnx-script"
}
functions {
  name: "aten_sigmoid"
  input: "self"
  output: "return_val"
  node {
    input: "self"
    output: "return_val"
    name: "n0"
    op_type: "Sigmoid"
    domain: ""
  }
  opset_import {
    domain: ""
    version: 18
  }
  domain: "onnx-script"
}

cc @gramalingam @liqunfu

[justinchuby]

test/onnx/test_fx_to_onnx_with_onnxruntime.py::TestFxToOnnxWithOnnxRuntime::test_sigmoid_add function: aten::add
function: aten::sigmoid
===========ONNX model: 
 ir_version: 7
producer_name: "pytorch"
producer_version: "2.0.0"
graph {
  node {
    output: "1"
    name: "Constant_0"
    op_type: "Constant"
    attribute {
      name: "value"
      t {
        data_type: 7
        raw_data: "\001\000\000\000\000\000\000\000"
      }
      type: TENSOR
    }
    doc_string: ""
  }
  node {
    input: "arg0_1"
    input: "1"
    output: "2"
    name: "aten_add_1"
    op_type: "aten_add"
    attribute {
      name: "alpha"
      i: 1
      type: INT
    }
    doc_string: ""
    domain: "onnx-script"
  }
  node {
    input: "2"
    output: "3"
    name: "aten_sigmoid_2"
    op_type: "aten_sigmoid"
    doc_string: ""
    domain: "onnx-script"
  }
  name: "torch_jit"
  input {
    name: "arg0_1"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 1
          }
          dim {
            dim_value: 4
          }
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 3
          }
        }
      }
    }
  }
  output {
    name: "3"
    type {
    }
  }
  value_info {
    name: "2"
    type {
    }
  }
}
opset_import {
  version: 14
}
opset_import {
  domain: "onnx-script"
  version: 1
}
functions {
  name: "aten_add"
  input: "self"
  input: "other"
  output: "return_val"
  attribute: "alpha"
  node {
    output: "alpha"
    name: "n0"
    op_type: "Constant"
    attribute {
      name: "value_float"
      type: FLOAT
      ref_attr_name: "alpha"
    }
    domain: ""
  }
  node {
    output: "tmp"
    name: "n1"
    op_type: "Constant"
    attribute {
      name: "value"
      t {
        data_type: 7
        int64_data: 1
        name: "tmp"
      }
      type: TENSOR
    }
    domain: ""
  }
  node {
    input: "alpha"
    input: "tmp"
    output: "tmp_0"
    name: "n2"
    op_type: "Equal"
    domain: ""
  }
  node {
    input: "tmp_0"
    output: "cond"
    name: "n3"
    op_type: "Not"
    domain: ""
  }
  node {
    input: "cond"
    output: "other_4"
    name: "n4"
    op_type: "If"
    attribute {
      name: "then_branch"
      g {
        node {
          output: "alpha_1"
          name: "n0"
          op_type: "Constant"
          attribute {
            name: "value_float"
            type: FLOAT
            ref_attr_name: "alpha"
          }
          domain: ""
        }
        node {
          input: "alpha_1"
          input: "other"
          output: "alpha_1_cast"
          name: "n1"
          op_type: "CastLike"
          domain: ""
        }
        node {
          input: "other"
          input: "alpha_1_cast"
          output: "other_2"
          name: "n2"
          op_type: "Mul"
          domain: ""
        }
        name: "thenGraph_4"
        output {
          name: "other_2"
          type {
          }
        }
      }
      type: GRAPH
    }
    attribute {
      name: "else_branch"
      g {
        node {
          input: "other"
          output: "other_3"
          name: "n0"
          op_type: "Identity"
          domain: ""
        }
        name: "elseGraph_4"
        output {
          name: "other_3"
          type {
          }
        }
      }
      type: GRAPH
    }
    domain: ""
  }
  node {
    input: "self"
    input: "other_4"
    output: "return_val"
    name: "n5"
    op_type: "Add"
    domain: ""
  }
  opset_import {
    domain: ""
    version: 18
  }
  domain: "onnx-script"
}
functions {
  name: "aten_sigmoid"
  input: "self"
  output: "return_val"
  node {
    input: "self"
    output: "return_val"
    name: "n0"
    op_type: "Sigmoid"
    domain: ""
  }
  opset_import {
    domain: ""
    version: 18
  }
  domain: "onnx-script"
}

===========ONNX model with inferred shapes: 
 ir_version: 7
producer_name: "pytorch"
producer_version: "2.0.0"
graph {
  node {
    output: "1"
    name: "Constant_0"
    op_type: "Constant"
    attribute {
      name: "value"
      t {
        data_type: 7
        raw_data: "\001\000\000\000\000\000\000\000"
      }
      type: TENSOR
    }
    doc_string: ""
  }
  node {
    input: "arg0_1"
    input: "1"
    output: "2"
    name: "aten_add_1"
    op_type: "aten_add"
    attribute {
      name: "alpha"
      i: 1
      type: INT
    }
    doc_string: ""
    domain: "onnx-script"
  }
  node {
    input: "2"
    output: "3"
    name: "aten_sigmoid_2"
    op_type: "aten_sigmoid"
    doc_string: ""
    domain: "onnx-script"
  }
  name: "torch_jit"
  input {
    name: "arg0_1"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 1
          }
          dim {
            dim_value: 4
          }
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 3
          }
        }
      }
    }
  }
  output {
    name: "3"
    type {
      tensor_type {
        elem_type: 1
      }
    }
  }
  value_info {
    name: "2"
    type {
      tensor_type {
        elem_type: 1
      }
    }
  }
  value_info {
    name: "1"
    type {
      tensor_type {
        elem_type: 7
        shape {
        }
      }
    }
  }
}
opset_import {
  version: 14
}
opset_import {
  domain: "onnx-script"
  version: 1
}
functions {
  name: "aten_add"
  input: "self"
  input: "other"
  output: "return_val"
  attribute: "alpha"
  node {
    output: "alpha"
    name: "n0"
    op_type: "Constant"
    attribute {
      name: "value_float"
      type: FLOAT
      ref_attr_name: "alpha"
    }
    domain: ""
  }
  node {
    output: "tmp"
    name: "n1"
    op_type: "Constant"
    attribute {
      name: "value"
      t {
        data_type: 7
        int64_data: 1
        name: "tmp"
      }
      type: TENSOR
    }
    domain: ""
  }
  node {
    input: "alpha"
    input: "tmp"
    output: "tmp_0"
    name: "n2"
    op_type: "Equal"
    domain: ""
  }
  node {
    input: "tmp_0"
    output: "cond"
    name: "n3"
    op_type: "Not"
    domain: ""
  }
  node {
    input: "cond"
    output: "other_4"
    name: "n4"
    op_type: "If"
    attribute {
      name: "then_branch"
      g {
        node {
          output: "alpha_1"
          name: "n0"
          op_type: "Constant"
          attribute {
            name: "value_float"
            type: FLOAT
            ref_attr_name: "alpha"
          }
          domain: ""
        }
        node {
          input: "alpha_1"
          input: "other"
          output: "alpha_1_cast"
          name: "n1"
          op_type: "CastLike"
          domain: ""
        }
        node {
          input: "other"
          input: "alpha_1_cast"
          output: "other_2"
          name: "n2"
          op_type: "Mul"
          domain: ""
        }
        name: "thenGraph_4"
        output {
          name: "other_2"
          type {
          }
        }
      }
      type: GRAPH
    }
    attribute {
      name: "else_branch"
      g {
        node {
          input: "other"
          output: "other_3"
          name: "n0"
          op_type: "Identity"
          domain: ""
        }
        name: "elseGraph_4"
        output {
          name: "other_3"
          type {
          }
        }
      }
      type: GRAPH
    }
    domain: ""
  }
  node {
    input: "self"
    input: "other_4"
    output: "return_val"
    name: "n5"
    op_type: "Add"
    domain: ""
  }
  opset_import {
    domain: ""
    version: 18
  }
  domain: "onnx-script"
}
functions {
  name: "aten_sigmoid"
  input: "self"
  output: "return_val"
  node {
    input: "self"
    output: "return_val"
    name: "n0"
    op_type: "Sigmoid"
    domain: ""
  }
  opset_import {
    domain: ""
    version: 18
  }
  domain: "onnx-script"
}

[justinchuby]

Another repro

import onnx
import onnxscript
from onnxscript import FLOAT, opset18
from onnxscript.function_libs.torch_aten.ops.core import aten_add, aten_sigmoid

"""
@onnxscript.script()
def aten_add(self: TReal, other: TReal, alpha: float = 1.0) -> TReal:
    # add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
    if alpha != 1.0:
        other = op.Mul(other, alpha)
    return op.Add(self, other)

@onnxscript.script()
def aten_sigmoid(self: TFloatOrBFloat16) -> TFloatOrBFloat16:
    # sigmoid(Tensor self) -> Tensor

    return op.Sigmoid(self)
"""

@onnxscript.script(default_opset=opset18)
def my_model(x: FLOAT[1, 2, 3, 4]):
    x = aten_add(x, 1.0, alpha=2.0)
    return aten_sigmoid(x)


p1 = my_model.to_model_proto()

print("before inference============")
print(p1)

p2 = onnx.shape_inference.infer_shapes(p1)

print("after inference============")
print(p2)
onnx.checker.check_model(p2)

Traceback (most recent call last):
  File "/home/justinchu/dev/onnx-script/test_onnx.py", line 22, in <module>
    onnx.checker.check_model(p2)
  File "/home/justinchu/anaconda3/envs/onnx/lib/python3.10/site-packages/onnx/checker.py", line 119, in check_model
    C.check_model(protobuf_string, full_check)
onnx.onnx_cpp2py_export.checker.ValidationError: Field 'shape' of 'type' is required but missing.

[jcwchen]

Thanks for the helpful reproduction steps. Actually onnx.onnx_cpp2py_export.checker.ValidationError: Field 'shape' of 'type' is required but missing. comes from model.graph.output. Graph's output requires "shape" to exist. A quick way to demonstrate:

import onnx
m = onnx.load("onnx_model2.onnx")
m.graph.output[0].type.tensor_type.shape.CopyFrom(onnx.TensorShapeProto())  # assign an empty shape
onnx.checker.check_model(m, full_check=True)

[gramalingam]

@justinchuby thanks!

I am not setup to try this out, but if you can, I suggest changing x = aten_add(x, 1, alpha=2.0) to x = aten_add(x, 1.0, alpha=2.0) ... not sure whether that will fix it, but the current inference introduces casts around constants when used in invocations of ops where we have an opschema. It doesn't work in function-calls like aten-add yet.

[justinchuby]

tried that too. This should work simply with onnxscript in the env

[gramalingam]

Hmmm. It seems to work (with the above change) for me, though I tried in 2 separate steps, so not exactly identical repro.

Another suggestion (general): I suggest changing onnx.shape_inference.infer_shapes(p1) to onnx.shape_inference.infer_shapes(p1, True, True) ... that enables more type-checking, which may be helpful.