Utils¶
A collection of useful functions for the kernel submodule.
-
get_op_name(tensor_name)[source]¶ Given a tensor name, return the corresponding op name.
- Parameters
tensor_name – The name of a tensor.
- Returns
The op name
- Return type
-
strip_replica_prefix(name)[source]¶ Given a tensor or op name, strip the AUTODIST-REPLICA prefix if there exists as the prefix.
- Parameters
name (string) – op or tensor name
- Returns
The stripped name
- Return type
-
parse_name_scope(name)[source]¶ Given a tensor or op name, return the name scope of the raw name.
- Parameters
name – The name of a tensor or an op.
- Returns
The name scope
- Return type
-
parse_optimizer_scope(update_op_name)[source]¶ Given the name of an update_op, return its optimizer name scope.
- Parameters
update_op_name – the name of an update_op (usually ResourceApply).
- Returns
The outermost name scope of the optimizer
- Return type
-
replica_prefix(replica_id)[source]¶ Generate replica prefix based on replica id.
Examples
>>> replica_prefix(1) {AUTODIST_REPLICA_PREFIX}1
-
get_consumers(op)[source]¶ Get a flat list from [output[0].consumers(), output[1].consumers(), …].
- Parameters
op – TensorFlow Operator
- Returns
The list of consumers
- Return type
List[Operation]
-
get_control_consumers(op)[source]¶ Get a flat list of the control-dependency consumers ([B]) of the op (A).
A: [B] A –> B B depends on A
- Parameters
op – TensorFlow Operator
- Returns
The list of control-dependency consumers
- Return type
List[Operation]
-
traverse(start_ops, end_ops=None, neighbors_fn=None)[source]¶ Traverse a graph and output the visited nodes.
- Parameters
start_ops (iter) – The nodes to start the traversal from.
end_ops (iter) – Optional. The nodes at which to stop traversing.
neighbors_fn (func) – Optional. Function from Op -> Iter[Op] that provides the neighbors of an op. Defaults to get_consumers.
- Returns
The visited nodes
- Return type
Set[Operation]
-
get_ancestors(start_ops, end_ops=None, include_control_inputs=False)[source]¶ Get all ancestor ops of the start ops.
Starting from start_ops, follow the computation graph backwards from consumer to input to find ancestors. Stop navigating the graph at end_ops. Include both start_ops and end_ops in the returning set of ancestor ops.
- Parameters
- Returns
The ancestor ops
- Return type
Set[Operation]
-
update_consumers(consumers, old_tensor, new_tensor)[source]¶ For each consumer’s inputs, replace old_tensor with new_tensor.
Be careful using op.consumers() directly as an argument, since this causes incorrect list iteration.
- Parameters
consumers (List[Operation]) – The consumer ops in the graph to be modified.
old_tensor (Tensor) – The tensor whose link to its consumer will be removed
new_tensor (Tensor) – The tensor which will replace old_tensor
-
update_control_consumers(control_consumer_ops, old_op, new_op)[source]¶ For each consumer’s control inputs, replace old_op with new_op.
- Parameters
control_consumer_ops (List[Operation]) – The control-dep consumer ops in the graph to be modified.
old_op (Operation) – The op whose link to its control-dep consumer will be removed
new_op (Operation) – The op which will replace old_op
-
update_colocation_group(ops, old_op, new_op)[source]¶ For each op in ops, we replace the colocation group as old_op to colocation group as new_op.
- Parameters
ops (Iterable[Operation]) – The operations to update
old_op (Operation) – The op having the old colocation group
new_op (Operation) – The op having the new colocation group
-
remove_from_control_consumers(control_consumer_ops, op_to_remove)[source]¶ Remove the op_to_remove from the control inputs for each op in “control_consumer_ops”.
- Parameters
control_consumer_ops (List[Operation]) – Ops that have op_to_remove as their current control inputs
op_to_remove (Operation) – The op to be removed