code: def test_setType_maintains_output_shape_for_single_custom_op_with_onnx_ops(
      self,
  ):
      self.addCleanup(
          torch.onnx.unregister_custom_op_symbolic, "::linalg_inv", 9
      )
  
      class CustomInverse(torch.nn.Module):
          def forward(self, x, y, z):
              x = torch.inverse(x)
              return x + y + z
  
      def linalg_inv_settype(g, self):
          return g.op("com.microsoft::Inverse", self).setType(
              self.type()
              .with_dtype(torch.float)
              .with_sizes([2, 3, 10, 10])
          )
  
      torch.onnx.register_custom_op_symbolic(
          "::linalg_inv", linalg_inv_settype, 9
      )
      model = CustomInverse()
      x = torch.randn(2, 3, 10, 10)
      y = torch.randn(2, 3, 10, 10)
      z = torch.randn(2, 3, 10, 10)
      f = io.BytesIO()
      torch.onnx.export(
          model,
          (x, y, z),
          f,
          opset_version=self.opset_version,
          custom_opsets={"com.microsoft": 1},
      )
  
      model_proto = onnx.load(io.BytesIO(f.getvalue()))
      # To validate the shape of inverse Op, we need to find inverse output name,
      # and then use it to identify its value_info for the shape.
      output_name = ""
      for node in model_proto.graph.node:
          if node.op_type == "Inverse":
              output_name = node.output[0]
              break
      assert output_name
      model_value_info = model_proto.graph.value_info
      self.assertIsNotNone(model_value_info)
      assert model_value_info
      for value_info in model_value_info:
          assert value_info.name
          if value_info.name == output_name:
              dims = value_info.type.tensor_type.shape.dim
              for i in range(len(dims)):
                  # If node output has shape info, it should have dim_value
                  # Otherwise, it has dim_params with dynamic shape
                  self.assertTrue(dims[i].HasField("dim_value"))
              for dim, rank in zip(dims, x.size()):
                  self.assertEqual(dim.dim_value, rank)
  
imports: import io
  import numpy
  import onnx
  import pytorch_test_common
  import torch
  from pytorch_test_common import skipIfUnsupportedMinOpsetVersion
  from torch.onnx import _constants, utils
  from torch.onnx._globals import GLOBALS
  from torch.onnx._internal import jit_utils
  from torch.testing._internal import common_utils
label_desc: Use image node_id 5 for calling the TestONNXCustomOpShapeInference obj's underlying member method code with example usage: obj.test_setType_maintains_output_shape_for_single_custom_op_with_onnx_ops() and returns: g, unknown