docstr: Overwrite TFConverter._get_stack() as TF2 generates different sub-graphs.
docstr_tok: ["Overwrite","TFConverter._get_stack","(",")","as","TF2","generates","different","sub-graphs","."]
code: def _get_stack(self, tfssa, root="main"):
      """
      Overwrite TFConverter._get_stack() as TF2 generates different sub-graphs.
      """
  
      # We're trying to get a order of how to loop through the graphs.
      # This is NOT necessarily a DAG.
      dep = {x: [] for x in tfssa.functions}
      for fname in tfssa.functions:
          for node in tfssa.functions[fname].graph.values():
              func_x, func_y = None, None
  
              if node.op in {"StatelessIf", "If"}:
                  func_x = node.attr.get("then_branch")
                  func_y = node.attr.get("else_branch")
              elif node.op in {"StatelessWhile", "While"}:
                  func_x = node.attr.get("body")
                  func_y = node.attr.get("cond")
  
              if func_x and fname not in dep[func_x]:
                  dep[func_x].append(fname)
              if func_y and fname not in dep[func_y]:
                  dep[func_y].append(fname)
  
      assert len(dep[root]) == 0
      graph_stack = simple_topsort(dep)
  
      return graph_stack
  
imports: from coremltools.converters.mil.frontend.tensorflow.basic_graph_ops import simple_topsort
  from coremltools.converters.mil.frontend.tensorflow.converter import TFConverter
label_desc: Use image node_id 1 for calling the TF2Converter obj's underlying member method code with example usage: obj._get_stack(tfssa, root) and returns: graph_stack