Home
last modified time | relevance | path

Searched refs:replicate_node (Results 1 – 2 of 2) sorted by relevance

/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
Ddistributed_tpu_rewrite_pass.h310 const Node* replicate_node, FunctionLibraryRuntime* flr,
340 const Node& replicate_node, const std::vector<InferredShape>& arg_shapes,
352 const Node* replicate_node, const NameAttrList& function,
370 Node* replicate_node, Node* compile_node,
430 int num_cores_per_replica, const Node& replicate_node,
515 const Node& replicate_node,
521 const DeviceSet& device_set, const Node& replicate_node,
529 int num_replicas, const Node& replicate_node, FunctionLibraryRuntime* flr,
536 const Node& replicate_node, const NameRangeMap& input_name_map,
544 const Node& replicate_node, Graph* graph,
[all …]
Ddistributed_tpu_rewrite_pass.cc569 const Node& replicate_node, in FillPaddingMap() argument
573 GetNodeAttr(replicate_node.attrs(), "padding_map", &padding_map_strs)); in FillPaddingMap()
585 Status GetStepMarkerLocation(const Node& replicate_node, in GetStepMarkerLocation() argument
588 TF_RETURN_IF_ERROR(GetNodeAttr(replicate_node.attrs(), "step_marker_location", in GetStepMarkerLocation()
1885 bool UseSpmdForXlaPartitioning(const Node* replicate_node) { in UseSpmdForXlaPartitioning() argument
1887 if (!replicate_node || in UseSpmdForXlaPartitioning()
1888 !TryGetNodeAttr(replicate_node->attrs(), "use_spmd_for_xla_partitioning", in UseSpmdForXlaPartitioning()
1917 const Node* replicate_node, FunctionLibraryRuntime* flr, in AssignArgsAndRetvalsToCores() argument
1939 for (const Edge* edge : replicate_node->out_edges()) { in AssignArgsAndRetvalsToCores()
1974 const bool use_spmd = (UseSpmdForXlaPartitioning(replicate_node) || in AssignArgsAndRetvalsToCores()
[all …]