code: def load(prog, **kwargs):
      if "main" not in prog.functions:
          msg = "main function not found in program {}"
          raise ValueError(msg.format(prog))
      if len(prog.functions) != 1:
          msg = (
              "Program must have exactly one `main` function to "
              "convert to NN. Program: {}"
          )
          raise ValueError(msg.format(prog))
  
      input_types = prog.main_input_types
      output_types = prog.main_output_types
  
      v1_inputs = []
      symbolic_inputs = {}
      for name, var in prog.functions["main"].inputs.items():
          if types.is_tensor(var.sym_type):
              sym_shape = var.sym_type.get_shape()
              if any_variadic(sym_shape):
                  raise NotImplementedError(
                      "Variadic rank is not supported"
                  )
              if any_symbolic(sym_shape):
                  user_specified = False
                  for input_type in input_types:
                      if name == input_type.name:
                          sym_shape = input_type.shape.default
                          user_specified = True
                          break
                  # Use dummy static shape, and will set it later.
                  shape = [
                      1 if is_symbolic(d) else d for d in sym_shape
                  ]
                  if not user_specified:
                      symbolic_inputs[name] = sym_shape
              else:
                  shape = sym_shape
              v1_inputs.append((name, Array(*shape)))
          elif types.is_scalar(var.sym_type):
              v1_inputs.append((name, Array(1)))
          else:
              raise NotImplementedError()
  
      v1_outputs = []
      for var in prog.functions["main"].outputs:
          if types.is_tensor(var.sym_type) or types.is_primitive(
              var.sym_type
          ):
              # Disregard the output types
              v1_outputs.append((var.name, None))
          else:
              raise NotImplementedError()
  
      # create neural network builder
      builder = neural_network.NeuralNetworkBuilder(
          v1_inputs,
          v1_outputs,
          disable_rank5_shape_mapping=True,
          use_float_arraytype=True,
      )
  
      # const in V2 are added lazily to V1 by each op whenever needed.
      # `const_context` stores the const names we've added so far and avoid
      # adding a const more than once.
      # const_context: list[set of str] (const name for v1 & v2
      # (the same)). Note that in NN in outer layer is visible from the inner
      # layer, so the const_context is simply a stack of set.
      const_context = []
      # Iterate through ops and add to builder
      convert_ops(
          const_context,
          builder,
          prog.functions["main"].operations,
          prog.functions["main"].outputs,
      )
  
      proto = builder.spec
      # image input
      has_image_input = any(
          [isinstance(s, ImageType) for s in input_types]
      )
      if has_image_input:
          proto = _convert_to_image_input(
              proto,
              input_types,
              skip_model_load=kwargs.get("skip_model_load", False),
          )
  
      # image output
      if output_types is not None:
          assert len(output_types) == len(
              prog.functions["main"].outputs
          ), "number of mil program outputs do not match the number of outputs provided by the user"
          for i, output_proto_desc in enumerate(
              proto.description.output
          ):
              output_var = prog.functions["main"].outputs[i]
              if isinstance(output_types[i], ImageType):
                  if not types.is_tensor(var.sym_type):
                      raise ValueError(
                          "Image output, '{}', is a scalar, but it should be a tensor of rank 4".format(
                              var.name
                          )
                      )
                  shape = var.sym_type.get_shape()
                  if any_variadic(shape):
                      raise ValueError(
                          "Variable rank model outputs, that are ImageTypes, are not supported"
                      )
                  if any([is_symbolic(d) for d in shape]):
                      raise NotImplementedError(
                          "Image output '{}' has symbolic dimensions in its shape".format(
                              var.name
                          )
                      )
                  _validate_image_input_output_shapes(
                      output_types[i].color_layout,
                      shape,
                      var.name,
                      is_input=False,
                  )
                  clr_space = _get_colorspace_enum(
                      output_types[i].color_layout
                  )
                  output_proto_desc.type.imageType.colorSpace = (
                      clr_space
                  )
                  output_proto_desc.type.imageType.width = shape[-1]
                  output_proto_desc.type.imageType.height = shape[-2]
  
      # classifier flag
      classifier_config = kwargs.get("classifier_config", None)
      if classifier_config is not None:
          # verify that classifier_config.predicted_probabilities_output if its exists.
          # And if its empty/None, fill it with the last non const op's output
          # this is done in "_get_probability_var_for_classifier()"
          probability_var = _get_probability_var_for_classifier(
              prog, classifier_config
          )
          if (
              classifier_config.predicted_probabilities_output
              != probability_var.name
          ):
              classifier_config.predicted_probabilities_output = (
                  probability_var.name
              )
          # add classifier related fields to the proto spec
          proto = _convert_to_classifier(
              proto,
              classifier_config,
              skip_model_load=kwargs.get("skip_model_load", False),
          )
  
      _set_user_inputs(proto, input_types)
      _set_symbolic_inputs(proto, symbolic_inputs)
      _set_optional_inputs(proto, input_types)
  
      return proto
  
file_path: coremltools/coremltools/converters/mil/backend/nn/load.py
imports: import coremltools
  from coremltools.converters._profile_utils import _profile
  from coremltools.converters.mil.backend.backend_helper import _get_probability_var_for_classifier
  from coremltools.converters.mil.input_types import ColorLayout, EnumeratedShapes, ImageType, RangeDim, Shape
  from coremltools.converters.mil.mil import types
  from coremltools.converters.mil.mil.types.symbolic import any_symbolic, any_variadic, is_symbolic
  from coremltools.models import MLModel
  from coremltools.models import neural_network
  from coremltools.models.datatypes import Array
  from coremltools.models.neural_network import flexible_shape_utils
  from coremltools.models.neural_network.flexible_shape_utils import add_enumerated_image_sizes, add_multiarray_ndshape_enumeration, set_multiarray_ndshape_range
  from ..backend_helper import _get_colorspace_enum, _validate_image_input_output_shapes
  from .op_mapping import convert_ops
label_desc: Use image node_id 6 for calling a global function with example usage: load(prog) and returns: proto