def set_graph_flow_context(graph: tf.Graph, active_context): """ sets graph context to active context provided :param graph: TensorFlow Graph (tf.Graph) :param active_context: context object to be set as current graph's context :return: """ # pylint: disable=protected-access graph._set_control_flow_context(active_context)
def updated_graph_flow_context_to_loop_context(graph: tf.Graph, preceeding_tensor: tf.Tensor): """ updates graph flow context to loop context :param graph: TensorFlow Graph (tf.Graph) :param preceeding_tensor: TF tensor that feeds into the op which needs modification :return: old graph context object """ # pylint: disable=protected-access old_graph_context = graph._get_control_flow_context() graph._set_control_flow_context( preceeding_tensor.op._get_control_flow_context()) return old_graph_context