Home
last modified time | relevance | path

Searched refs:computation_map (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dsharding_propagation.cc285 const ComputationMap& computation_map, in SupportSpatialPartitioning() argument
288 computation_map.find(instruction->parent()) == computation_map.end()) { in SupportSpatialPartitioning()
327 return computation_map.find(instruction->parent()) != in SupportSpatialPartitioning()
328 computation_map.end(); in SupportSpatialPartitioning()
578 const ComputationMap& computation_map, in InferShardingFromOperands() argument
602 if (!SupportSpatialPartitioning(instruction, computation_map, is_spmd)) { in InferShardingFromOperands()
863 auto parent_it = computation_map.find(instruction->parent()); in InferShardingFromOperands()
864 if (parent_it == computation_map.end()) { in InferShardingFromOperands()
1408 const ComputationMap& computation_map, in InferShardingFromUsers() argument
1428 if (!SupportSpatialPartitioning(instruction, computation_map, is_spmd)) { in InferShardingFromUsers()
[all …]
Dhlo_module.cc363 absl::flat_hash_map<int64, HloComputation*> computation_map; in CreateFromProto() local
370 HloComputation::CreateFromProto(computation_proto, computation_map, in CreateFromProto()
375 TF_RET_CHECK(!ContainsKey(computation_map, computation_id)); in CreateFromProto()
376 computation_map[computation_id] = computation.get(); in CreateFromProto()
Dhlo_instruction.cc70 const absl::flat_hash_map<int64, HloComputation*>& computation_map, in CreateFromProto() argument
116 const auto computations = [&computation_map, &proto](int index) { in CreateFromProto()
117 return computation_map.at(proto.called_computation_ids(index)); in CreateFromProto()
119 const auto all_computations = [&computation_map, &proto]() { in CreateFromProto()
123 [&computation_map](int64 computation_id) { in CreateFromProto()
124 return computation_map.at(computation_id); in CreateFromProto()
136 [&](int64 id) { return computation_map.contains(id); })) in CreateFromProto()
362 tensorflow::gtl::FindPtrOrNull(computation_map, fusion_id); in CreateFromProto()
779 computation_map.at(computation_id)); in CreateFromProto()
Dhlo_computation.h236 const absl::flat_hash_map<int64, HloComputation*>& computation_map,
Dhlo_computation.cc664 const absl::flat_hash_map<int64, HloComputation*>& computation_map, in CreateFromProto() argument
673 instruction_proto, instruction_map, computation_map, in CreateFromProto()
Dhlo_instruction.h500 const absl::flat_hash_map<int64, HloComputation*>& computation_map,