def _set_read_only_resource_inputs_attr(op, func_graph): """Sets the list of resource inputs which are read-only. This is used by AutomaticControlDependencies. Args: op: PartitionedCall Operation. func_graph: FuncGraph. """ read_only_indices = acd.get_read_only_resource_input_indices_graph(func_graph) ops.set_int_list_attr(op, acd.READ_ONLY_RESOURCE_INPUTS_ATTR, read_only_indices)
def _set_read_only_resource_inputs_attr(op, branch_graphs): """Sets the list of resource inputs which are read-only. This is used by AutomaticControlDependencies. Args: op: If or Case Operation. branch_graphs: List of branch FuncGraphs. """ # The first entry in `op.inputs` is the predicate which is not passed to # branch graphs so len(branch_graph[i].inputs) == len(op.inputs) - 1. read_only_indices = set(range(len(op.inputs) - 1)) for branch_graph in branch_graphs: assert len(branch_graph.inputs) == len(op.inputs) - 1, "should never happen" if not read_only_indices: break branch_read_only_indices = acd.get_read_only_resource_input_indices_graph( branch_graph) read_only_indices = read_only_indices.intersection(branch_read_only_indices) # Convert indices in `branch_graphs[i].inputs` to `op.inputs`. read_only_indices = [i + 1 for i in read_only_indices] ops.set_int_list_attr(op, acd.READ_ONLY_RESOURCE_INPUTS_ATTR, sorted(read_only_indices))