code: """Converts SavedModel into an optimized protobuf for inference"""
  
  from dataclasses import dataclass, field
  import logging
  import re
  from typing import List, Dict, Iterable, Union
  
  import tensorflow as tf
  from tensorflow.python.tools.freeze_graph import freeze_graph_with_def_protos
  from tensorflow.python.framework.graph_util import extract_sub_graph
  
  from deepr.io import Path
  from deepr.jobs import base
  from deepr.utils.graph import INIT_ALL_TABLES
  
  
  LOGGER = logging.getLogger(__name__)
  
  
  class TensorsNotFoundError(ValueError):
      """No tensors found for those names.
  
      Most Tensorflow operators take a ``name`` argument that should be
      used if you want to use a custom name for a tensor, otherwise a
      default will be used.
  
      If accessing the operator creating a Tensor is not possible, you
      can use ``tf.identity`` to name the Tensor. However, note that it
      adds a new op to the graph that creates a copy of the input Tensor,
      and thus should be limited to avoid overhead.
      """
  
      def __init__(self, tensors: Iterable[str]):
          msg = f"Tensors not found: {tensors}. Use `tf.identity(t, name)` to name nodes in your graph."
          super().__init__(msg)
  
  
  @dataclass
  class OptimizeSavedModel(base.Job):
      """Converts SavedModel into an optimized protobuf for inference
  
      This job reads the input SavedModel, rename some nodes using the
      ``new_names`` argument (raises an error if some renames cannot be
      found), create placeholders given by ``feeds`` (and removes all
      other placeholders not in this list), and finally freezes the sub
      graph that produces the output tensor ``fetch``.
  
      When creating the original SavedModel, it is recommended to use
      ``tf.identity`` operators to mark some tensors as future feeds or
      fetches.
  
      WARNING: successful completion of this job is no guarantee that the
      exported Graph is correct. It is recommended to test the export in
      a separate job.
  
      Attributes
      ----------
      path_saved_model : str
          Path to directory containing SavedModel exports to convert
      path_optimized_model : str
          Path to directory that will contain the export
      graph_name : str
          Name of the output graph (name of the protobuf file)
      new_names : Dict[str, str]
          Mapping old names (SavedModel nodes) -> new names (export)
      blacklisted_variables : Union[str, List[str]]
          List of variable names not to include in the export
      feeds : Union[str, List[str]]
          List of nodes to use as inputs, or comma separated string.
      fetch : Union[str, List[str]]
          List of nodes to use as output, or comma separated string.
      """
  
      path_saved_model: str
      path_optimized_model: str
      graph_name: str
      feeds: Union[str, List[str]]
      fetch: Union[str, List[str]]
      new_names: Dict[str, str] = field(default_factory=dict)
      blacklisted_variables: List[str] = field(default_factory=list)
  
      def run(self):
          # Normalize feeds and fetch
          fetch = self.fetch.split(",") if isinstance(self.fetch, str) else self.fetch
          feeds = self.feeds.split(",") if isinstance(self.feeds, str) else self.feeds
  
          # Find latest SavedModel export in path_saved_model
          subdirs = [
              str(path) for path in Path(self.path_saved_model).iterdir() if path.is_dir() and "temp" not in str(path)
          ]
          latest = str(sorted(subdirs)[-1])
          LOGGER.info(f"Using SavedModel {latest}")
  
          # Reload SavedModel Graph, optimize and export
          with tf.compat.v1.Session(graph=tf.Graph()) as sess:
              meta_graph_def = tf.compat.v1.saved_model.loader.load(sess, ["serve"], latest)
              graph_def = meta_graph_def.graph_def
  
              # Add table initializer if present, or create it
              if INIT_ALL_TABLES in {node.name for node in graph_def.node}:
                  fetch.append(INIT_ALL_TABLES)
              else:
                  table_initializers = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TABLE_INITIALIZERS)
                  if table_initializers:
                      LOGGER.info(f"Adding {INIT_ALL_TABLES} Node to the graph")
                      table_init_op = tf.group(*table_initializers, name=INIT_ALL_TABLES)
                      node_def = table_init_op.node_def
                      graph_def.node.append(node_def)
                      fetch.append(INIT_ALL_TABLES)
  
              # Rename nodes
              graph_def = rename_nodes(graph_def, self.new_names)
  
              # Setup (create / remove) placeholders
              graph_def = make_placeholders(graph_def, feeds)
  
              # Keep only part of the graph that produces tensor 'fetch'
              graph_def = extract_sub_graph(graph_def, fetch)
  
              # Replace variables by constants
              graph_def = freeze_graph_with_def_protos(
                  input_graph_def=graph_def,
                  input_saver_def=None,
                  input_checkpoint=None,
                  output_node_names=",".join(fetch),
                  restore_op_name=None,
                  filename_tensor_name=None,
                  output_graph=None,
                  clear_devices=True,
                  initializer_nodes=None,
                  variable_names_blacklist=",".join(self.blacklisted_variables),
                  input_saved_model_dir=latest,
                  saved_model_tags=["serve"],
              )
              tf.io.write_graph(graph_def, logdir=self.path_optimized_model, name=self.graph_name, as_text=False)
              LOGGER.info(f"Optimized Model successfully exported to {self.path_optimized_model}/{self.graph_name}")
  
  
  def rename_nodes(graph_def: tf.compat.v1.GraphDef, new_names: Dict[str, str]) -> tf.compat.v1.GraphDef:
      """Rename items in the graph to new ones defined in new_names
  
      Parameters
      ----------
      graph_def : tf.GraphDef
          Graph Definition
      new_names : Dict[str, str]
          Mapping old name -> new name
  
      Returns
      -------
      tf.GraphDef
          A copy of the input GraphDef with renamed nodes
  
      Raises
      ------
      TensorsNotFoundError
          If new_names refers to an node not found in graph_def
      """
      # Create copy of each node with a new name
      nodes = []
      for node in graph_def.node:
          new_node = tf.compat.v1.NodeDef()
          new_node.CopyFrom(node)
          nodes.append(new_node)
          if node.name in new_names:
              new_node.name = new_names[node.name]
              LOGGER.info(f"Node renamed: {node.name} -> {new_node.name}")
  
      # Check that all new names were used
      if not set(new_names.values()) <= set(node.name for node in nodes):
          missing = set(new_names.values()) - set(node.name for node in nodes)
          raise TensorsNotFoundError(missing)
  
      # Update node references (inputs and location) to renamed nodes
      for node in nodes:
          for idx, name in enumerate(node.input):
              node.input[idx] = new_names[name] if name in new_names else name
          if "_class" in node.attr:
              attr = node.attr["_class"]
              for idx, item in enumerate(attr.list.s):
                  loc_match = re.match(r"^loc:@(.+)$", item.decode())
                  if loc_match and loc_match.groups()[0] in new_names:
                      new_name = new_names[loc_match.groups()[0]]
                      attr.list.s[idx] = f"loc:@{new_name}".encode()
  
      # Create Graph with renamed nodes
      new_graph = tf.compat.v1.GraphDef()
      new_graph.node.extend(nodes)
      return new_graph
  
  
  def make_placeholders(graph_def: tf.compat.v1.GraphDef, names: List[str]) -> tf.compat.v1.GraphDef:
      """Create placeholders for names and remove other placeholders
  
      Parameters
      ----------
      graph_def : tf.GraphDef
          Graph definition
      names : List[str]
          Names of placeholders to keep / create for this graph
  
      Returns
      -------
      tf.GraphDef
          A copy of the input GraphDef with new placeholders
  
      Raises
      ------
      TensorsNotFoundError
          If names refers to a node that is not present
      """
      # Create copy of each node and change to Placeholder if in names
      nodes = []
      for node in graph_def.node:
          if node.name not in names and node.op == "Placeholder":
              LOGGER.info(f"Removing placeholder {node.name}")
              continue
          new_node = tf.compat.v1.NodeDef()
          if node.name in names and node.op != "Placeholder":
              LOGGER.info(f"Creating placeholder {node.name}")
              new_node.name = node.name
              new_node.op = "Placeholder"
              new_node.attr["shape"].CopyFrom(tf.compat.v1.AttrValue(shape=node.attr["_output_shapes"].list.shape[0]))
              new_node.attr["dtype"].CopyFrom(node.attr["T"])
          else:
              new_node.CopyFrom(node)
          nodes.append(new_node)
  
      # Check that all expected placeholders have been found
      if not set(names) <= set(node.name for node in nodes):
          missing = set(names) - set(node.name for node in nodes)
          raise TensorsNotFoundError(missing)
  
      # Create Graph with renamed nodes
      new_graph = tf.compat.v1.GraphDef()
      new_graph.node.extend(nodes)
      return new_graph
  
apis: tensorflow.io.write_graph tensorflow.python.framework.graph_util.extract_sub_graph tensorflow.compat.v1.NodeDef tensorflow.compat.v1.AttrValue tensorflow.Graph tensorflow.compat.v1.GraphDef tensorflow.compat.v1.saved_model.loader.load tensorflow.group tensorflow.compat.v1.get_collection