support for subfunctions included in a subgraphs

onnxscript/converter.py

@@ -87,6 +87,7 @@ def _known_modules():
     return {
         'onnxscript': onnxscript,
         'onnxscript.onnx': onnxscript.onnx,
+        'onnxscript.values': onnxscript.values,
         'onnxscript.onnx_types': onnxscript.onnx_types,
         'onnxscript.onnx.opset15': onnxscript.onnx.opset15
     }
@@ -587,7 +588,7 @@ def translate_function_def(self, fn: ast.FunctionDef):
             warn(f"{fn.name}: Default values not yet implemented.")
         if args.vararg or args.kwonlyargs or args.kw_defaults or args.kwarg:
             warn(f"{fn.name}: Unsupported feature in function signature.")
-        domain = self.globals["__opset_domain__"] if "__opset_domain__" in self.globals else ""
+        domain = self.globals["__opset_domain__"] if "__opset_domain__" in self.globals else "this"
         self.current_fn = self.ir_builder.new_function(fn.name, domain)
         for x in args.args:
             if x.annotation:

onnxscript/test/converter_test.py

@@ -8,6 +8,7 @@
 from onnx.helper import printable_graph
 from onnx.onnx_cpp2py_export.checker import ValidationError
 import onnxruntime
+from onnxruntime.capi.onnxruntime_pybind11_state import Fail
 from onnxscript.converter import Converter
 from onnxscript.values import Opset
 
@@ -29,23 +30,26 @@ def _convert_and_save(self, script, save_text=False, check_ort=False):
             with self.subTest(f=f.name):
                 model = f.to_model_proto(producer_name='p2o')
                 if save_text:
-                    with open(os.path.join(TEST_OUTPUT_DIR, f.name + ".txt"), 'w') as f:
-                        f.write(printable_graph(model.graph))
+                    with open(os.path.join(TEST_OUTPUT_DIR, f.name + ".txt"), 'w') as fi:
+                        fi.write(printable_graph(model.graph))
                         for fct in model.functions:
-                            f.write("\n-------------------------\n")
-                            f.write(printable_graph(fct))
+                            fi.write("\n-------------------------\n")
+                            fi.write(printable_graph(fct))
                 if check_ort:
-                    onnxruntime.InferenceSession(model.SerializeToString())
+                    try:
+                        onnxruntime.InferenceSession(model.SerializeToString())
+                    except Fail as e:
+                        raise AssertionError(f"onnxruntime cannot load function {f.name}\n{str(model)}") from e
                 model = onnx.shape_inference.infer_shapes(model)
                 if save_text:
-                    with open(os.path.join(TEST_OUTPUT_DIR, f.name + ".shape.txt"), 'w') as f:
-                        f.write(printable_graph(model.graph))
+                    with open(os.path.join(TEST_OUTPUT_DIR, f.name + ".shape.txt"), 'w') as fi:
+                        fi.write(printable_graph(model.graph))
                         for fct in model.functions:
                             f.write("\n-------------------------\n")
                             f.write(printable_graph(fct))
                 try:
                     onnx.checker.check_model(model)
-                except ValidationError as e:
+                except (ValidationError, AssertionError) as e:
                     onnx.save(model, os.path.join(TEST_OUTPUT_DIR, f.name + ".error.onnx"))
                     raise AssertionError(
                         "Verification of model failed.") from e
@@ -101,12 +105,15 @@ def test_onnxfns1A(self):
     def test_models(self):
         self._convert_and_save(os.path.join(TEST_INPUT_DIR, "onnxmodels.py"))
 
-    def test_subfunction(self):
-        from .models import subfunction
+    def test_subfunction_check_model(self):
+        from onnxscript.test.models import subfunction
         model = subfunction.MyElu.function_ir.to_model_proto(producer_name='p2o')
         model = onnx.shape_inference.infer_shapes(model)
         onnx.checker.check_model(model)
 
+    def test_subfunction(self):
+        self._convert_and_save(os.path.join(TEST_INPUT_DIR, "subfunction.py"), check_ort=True)
+
     def test_if_models(self):
         self._convert_and_save(os.path.join(TEST_INPUT_DIR, "if_statement.py"))
 
@@ -123,4 +130,5 @@ def test_docstring(self):
 if __name__ == '__main__':
     # import logging
     # logging.basicConfig(level=logging.DEBUG)
+    TestConverter().test_subfunction()
     unittest.main()

onnxscript/test/models/subfunction.py

@@ -5,27 +5,35 @@
 from onnxscript.values import CustomOpset
 from onnxscript.onnx import opset15 as op
 
-opset = CustomOpset('this', 1)
 
-@script(opset)
+@script(CustomOpset('this', 1))
 def MySelu(X: FLOAT[None], alpha: FLOAT[1], gamma: FLOAT[1]) -> FLOAT[None]:
     neg = gamma * (alpha * op.Exp(X) - alpha)
     pos = gamma * X
-    return op.Where(X <= 0, neg, pos)
+    return op.Where(X <= 0., neg, pos)
 
-@script(opset)
+@script(CustomOpset('this', 1))
 def MyElu(X: FLOAT[None], beta: FLOAT[1]) -> FLOAT[None]:
     alpha = op.Constant(value_float=1.)
     return MySelu(X, alpha, beta)
 
-@script(opset)
+@script(CustomOpset('this', 1))
 def MyEluB(X: FLOAT[None], beta: FLOAT[1]) -> FLOAT[None]:
     alpha = op.Constant(value_float=1.)
     res = MySelu(X, alpha, beta)
     return res
 
-@script(opset)
+@script(CustomOpset('this', 1))
 def MyEluC(X: FLOAT[None], beta: FLOAT[1]) -> FLOAT[None]:
     alpha = op.Constant(value_float=1.)
     res = op.Identity(MySelu(X, alpha, beta))
     return res
+
+@script(CustomOpset('this', 1))
+def IfMyEluD(X: FLOAT[None], beta: FLOAT[1]) -> FLOAT[None]:
+    zero = op.Constant(value_float=1.)
+    if beta > 0:
+        result = MyEluB(X, beta)
+    else:
+        result = MyEluC(X, beta)
+    return result