def _search(g, node_name):
        """
        Find the nearest Switch nodes upstream of node_name.
        """
        node = g[node_name]

        switches = (
            FindAllUpstreamTerminals(lambda x: x.op == "Switch")
            .visit(g, node.name)
            .get_result()
        )
        if len(switches) == 0:
            switches = (
                FindAllUpstreamTerminals(
                    lambda x: x.op == "Switch" or x.attr.get("was_switch") is not None
                )
                .visit(g, node.name)
                .get_result()
            )
        return switches

        Find the nearest Switch nodes upstream of node_name.
        
coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/cond_to_where.py