Searched refs:index_mapping (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | rearrange_function_argument.cc | 42 const std::vector<int>& index_mapping) { in ShuffleInputDataTypeAttribute() argument
43 std::vector<DataType> result(index_mapping.size()); in ShuffleInputDataTypeAttribute()
45 result[index_mapping.at(i)] = in_types[i]; in ShuffleInputDataTypeAttribute()
57 std::vector<int>* index_mapping) { in InputTypesNeedsRearrange() argument
92 index_mapping->resize(in_types.size()); in InputTypesNeedsRearrange()
95 (*index_mapping)[i] = non_resource_index; in InputTypesNeedsRearrange()
98 (*index_mapping)[i] = resource_index; in InputTypesNeedsRearrange()
109 const std::vector<int>& index_mapping) { in ReorderInputEdges() argument
121 int new_dst_input = index_mapping.at(dst_input); in ReorderInputEdges()
134 const std::vector<int>& index_mapping) { in ReorderOutputEdges() argument
[all …]
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | encapsulate_tpu_computations_pass.cc | 1945 const absl::flat_hash_map<int, int>& index_mapping, in RemoveArgsToLiftFromFunctionBody() argument
1951 int new_index = index_mapping.at(i); in RemoveArgsToLiftFromFunctionBody()
1997 Status CleanUpInEdges(const absl::flat_hash_map<int, int>& index_mapping, in CleanUpInEdges() argument
2008 if (!index_mapping.contains(arg_idx)) { in CleanUpInEdges()
2014 if (index_mapping.at(arg_idx) == arg_idx) continue; in CleanUpInEdges()
2017 index_mapping.at(arg_idx) + arg_to_input_edge_offset); in CleanUpInEdges()
2024 Status UpdateTypeAttribute(const absl::flat_hash_map<int, int>& index_mapping, in UpdateTypeAttribute() argument
2028 new_dtypes.reserve(index_mapping.size()); in UpdateTypeAttribute()
2030 if (index_mapping.contains(i)) { in UpdateTypeAttribute()
2093 absl::flat_hash_map<int, int> index_mapping; in ArgIndexMapping() local
[all …]
|
| D | distributed_tpu_rewrite_pass.cc | 2812 std::vector<NodeOut> index_mapping; in CreateOrGetPerHostVariableCopy() local
2813 index_mapping.reserve(variable_reads.size()); in CreateOrGetPerHostVariableCopy()
2821 index_mapping.push_back( in CreateOrGetPerHostVariableCopy()
2826 index_mapping.push_back(NodeOut{read, 0}); in CreateOrGetPerHostVariableCopy()
2843 if (index_mapping[i].node == nullptr) { in CreateOrGetPerHostVariableCopy()
2845 index_mapping[i].node = id_node; in CreateOrGetPerHostVariableCopy()
2852 graph->AddEdge(variable_reads[i], 0, id_node, index_mapping[i].index); in CreateOrGetPerHostVariableCopy()
2864 graph->AddEdge(variable_reads[i], 0, id_node, index_mapping[i].index); in CreateOrGetPerHostVariableCopy()
2874 graph->AddEdge(var_read, 0, id_node, index_mapping[i].index); in CreateOrGetPerHostVariableCopy()
2879 auto result = index_mapping[var_index]; in CreateOrGetPerHostVariableCopy()
[all …]
|