Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/dtensor/mlir/
Dmove_compilation_to_host.cc212 std::map<Mesh, mlir::TF::StatefulPartitionedCallOp>& computation_map, in HandleCompilationOps() argument
216 computation_map, [](const auto& it) { return it.first.is_tpu_mesh(); }); in HandleCompilationOps()
218 if (xla_mesh == computation_map.end()) { in HandleCompilationOps()
232 auto mesh_it = llvm::find_if(computation_map, [&](auto& it) { in HandleCompilationOps()
235 if (mesh_it == computation_map.end()) in HandleCompilationOps()
382 std::map<Mesh, mlir::TF::StatefulPartitionedCallOp> computation_map; in runOnOperation() local
384 IdentifyAndValidateMeshComputations(main_func, &computation_map))) in runOnOperation()
388 if (mlir::failed(HandleCompilationOps(compilation_key_ops, computation_map, in runOnOperation()
/external/tensorflow/tensorflow/compiler/xla/service/
Dsharding_propagation.cc323 const ShardingPropagation::ComputationMap& computation_map, bool is_spmd, in SupportSpatialPartitioning() argument
331 computation_map.find(instruction->parent()) == computation_map.end() & in SupportSpatialPartitioning()
374 return computation_map.find(instruction->parent()) != in SupportSpatialPartitioning()
375 computation_map.end(); in SupportSpatialPartitioning()
719 const ShardingPropagation::ComputationMap& computation_map, in InferShardingFromUsers() argument
749 if (!SupportSpatialPartitioning(instruction, computation_map, is_spmd, false, in InferShardingFromUsers()
1707 HloInstruction* instruction, const ComputationMap& computation_map, in InferShardingFromOperands() argument
1738 if (!SupportSpatialPartitioning(instruction, computation_map, is_spmd_, in InferShardingFromOperands()
2083 auto parent_it = computation_map.find(instruction->parent()); in InferShardingFromOperands()
2084 if (parent_it == computation_map.end()) { in InferShardingFromOperands()
[all …]
Dsharding_propagation.h89 const ComputationMap& computation_map,
Dhlo_module.cc422 absl::flat_hash_map<int64_t, HloComputation*> computation_map; in CreateFromProto() local
429 HloComputation::CreateFromProto(computation_proto, computation_map, in CreateFromProto()
434 TF_RET_CHECK(!ContainsKey(computation_map, computation_id)); in CreateFromProto()
435 computation_map[computation_id] = computation.get(); in CreateFromProto()
Dhlo_instruction.cc87 const absl::flat_hash_map<int64_t, HloComputation*>& computation_map, in CreateFromProto() argument
133 const auto computations = [&computation_map, &proto](int index) { in CreateFromProto()
134 return computation_map.at(proto.called_computation_ids(index)); in CreateFromProto()
136 const auto all_computations = [&computation_map, &proto]() { in CreateFromProto()
140 [&computation_map](int64_t computation_id) { in CreateFromProto()
141 return computation_map.at(computation_id); in CreateFromProto()
153 [&](int64_t id) { return computation_map.contains(id); })) in CreateFromProto()
419 tensorflow::gtl::FindPtrOrNull(computation_map, fusion_id); in CreateFromProto()
963 computation_map.at(computation_id)); in CreateFromProto()
Dhlo_computation.h275 const absl::flat_hash_map<int64_t, HloComputation*>& computation_map,
Dhlo_computation.cc657 const absl::flat_hash_map<int64_t, HloComputation*>& computation_map, in CreateFromProto() argument
666 instruction_proto, instruction_map, computation_map, in CreateFromProto()
Dhlo_instruction.h533 const absl::flat_hash_map<int64_t, HloComputation*>& computation_map = {},
/external/tensorflow/tensorflow/compiler/xla/client/
Dvalue_inference.cc171 absl::flat_hash_map<int64_t, HloComputation*> computation_map; in Evaluate() local
176 computation_map[inst.called_computation_ids(0)] = computation; in Evaluate()
180 HloInstruction::CreateFromProto(inst, operand_map, computation_map)); in Evaluate()