def _populate_sub_graph_input_shapes(self, graph, graph_fns):
        """
        Populate function (sub-graph) input shapes from control flow op's inputs
        Note that the functions (sub-graphs) are not nested but the control flow
        ops are nested. The input shapes are used to extract sub-graphs from the
        parent graph (as the input of function_def_to_graph).

        Parameter
        ---------
        graph: tf.Graph
            TensorFlow graph.
        graph_fns: list of graph functions.
            List of TensorFlow graph functions.

        Returns
        -------
        sg_input_shapes: dict(str: list)
            Dictionary of function (sub-graph) name and input shape pairs.
        """
        sg_input_shapes = {}
        sub_graphs = []
        for op in graph.get_operations():
            if op.type not in {"StatelessIf", "If", "StatelessWhile", "While"}:
                continue

            sg1, sg2 = None, None
            if op.type in {"StatelessIf", "If"}:
                sg1 = op.get_attr("then_branch").name
                sg2 = op.get_attr("else_branch").name
            if op.type in {"StatelessWhile", "While"}:
                sg1 = op.get_attr("cond").name
                sg2 = op.get_attr("body").name

            # memorize input shapes for sub-graph conversions
            op_input_shapes = [i.get_shape() for i in op.inputs]
            sg_input_shapes.update({sg1: op_input_shapes, sg2: op_input_shapes})
            sub_graphs += [sg1, sg2]

        for name in sub_graphs:
            sg = graph_fns.get(name)
            fn_def = context.get_function_def(name)
            op_input_shapes = sg_input_shapes[name]
            op_input_shapes = op_input_shapes[-len(fn_def.signature.input_arg) :]
            fn_graph = _function_def_to_graph(fn_def, input_shapes=op_input_shapes)
            sg_input_shapes.update(
                self._populate_sub_graph_input_shapes(fn_graph, graph_fns)
            )

        return sg_input_shapes

        Populate function (sub-graph) input shapes from control flow op's inputs
        Note that the functions (sub-graphs) are not nested but the control flow
        ops are nested. The input shapes are used to extract sub-graphs from the
        parent graph (as the input of function_def_to_graph).

        Parameter
        ---------
        graph: tf.Graph
            TensorFlow graph.
        graph_fns: list of graph functions.
            List of TensorFlow graph functions.

        Returns
        -------
        sg_input_shapes: dict(str: list)
            Dictionary of function (sub-graph) name and input shape pairs.
        
coremltools/converters/mil/frontend/tensorflow2/load.py