Rename intermediate values to _val and _const | feat(torchlib)

[ghstack-poisoned]

Author: justinchuby

onnxscript/function_libs/torch_lib/graph_building.py

@@ -76,6 +76,18 @@
 # TODO(justinchuby): Build a context manager to handle source information.
 
 
+def _rename_intermediate_value(name: str) -> str:
+    if name.isdigit():
+        return f"_val_{name}"
+    return name
+
+
+def _rename_intermediate_constant(name: str) -> str:
+    if name.isdigit():
+        return f"_const_{name}"
+    return name
+
+
 class TorchScriptTensor(onnxscript_tensor.Tensor):
     """A onnxscript tensor that wraps a torchscript Value."""
 
@@ -454,6 +466,7 @@ def _add_constant_to_graph(self, constant) -> torch.Value:
                 self._torch_graph, "prim::Constant", inputs=(), attributes={}
             )[0]
             value.setType(torch.OptionalType.ofTensor())
+            value.setDebugName(_rename_intermediate_constant(value.debugName()))
             return value
 
         if isinstance(constant, bool):
@@ -475,12 +488,14 @@ def _add_constant_to_graph(self, constant) -> torch.Value:
             raise TypeError(
                 f"Constant input '{constant}' of type '{type(constant)}' is not supported"
             )
-        return _create_op_call_in_torch_graph(
+        value = _create_op_call_in_torch_graph(
             self._torch_graph,
             "onnx::Constant",
             inputs=(),
             attributes=dict(value=constant_tensor),
         )[0]
+        value.setDebugName(_rename_intermediate_constant(value.debugName()))
+        return value
 
     @runtime_typing.checked
     def _add_torchscript_op_call(
@@ -524,9 +539,15 @@ def _add_torchscript_op_call(
             attributes=onnx_attributes,
             n_outputs=n_outputs,
         )
-        if len(result) <= 1:
-            return TorchScriptTensor(result[0])
-        return tuple(TorchScriptTensor(v) for v in result)
+        assert result, "Expected at least one output from ONNX op call."
+        if len(result) == 1:
+            tensor = TorchScriptTensor(result[0])
+            tensor.name = _rename_intermediate_value(tensor.name)
+            return tensor
+        tensors = tuple(TorchScriptTensor(v) for v in result)
+        for tensor in tensors:
+            tensor.name = _rename_intermediate_value(tensor.name)
+        return tensors
 
     @runtime_typing.checked
     def fetch_function_proto_dict(