Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dflatten_call_graph.cc35 HloComputation* new_computation) { in ReplaceCalledComputation() argument
39 instruction->set_while_condition(new_computation); in ReplaceCalledComputation()
42 instruction->set_while_body(new_computation); in ReplaceCalledComputation()
48 instruction->set_to_apply(new_computation); in ReplaceCalledComputation()
57 instruction->set_branch_computation(b, new_computation); in ReplaceCalledComputation()
68 [&](HloComputation*) { return new_computation; }); in ReplaceCalledComputation()
69 new_computation->AddAsyncInstruction(instruction); in ReplaceCalledComputation()
H A Dhlo_module.cc206 HloComputation* new_computation = tensorflow::gtl::FindWithDefault( in ReplaceComputations() local
208 if (new_computation != nullptr) { in ReplaceComputations()
209 instruction->set_branch_computation(b, new_computation); in ReplaceComputations()
884 HloComputation* new_computation = context.FindComputation(computation); in Clone() local
888 if (new_computation != nullptr) { in Clone()
890 clone_schedule.GetOrCreateSequence(new_computation); in Clone()
945 HloComputation* new_computation; in DeepCloneComputation() local
947 if ((new_computation = context->FindComputation(computation)) != nullptr) { in DeepCloneComputation()
948 return new_computation; in DeepCloneComputation()
950 new_computation = in DeepCloneComputation()
[all …]
H A Drng_bit_generator_expander.cc95 HloComputation* new_computation = in GetGeneratorComputation() local
97 computation_cache_.emplace(cache_key, new_computation); in GetGeneratorComputation()
98 return new_computation; in GetGeneratorComputation()
H A Dhlo_clone_context.h50 HloComputation* new_computation) { in MapComputation() argument
51 computations_[old_computation] = new_computation; in MapComputation()
H A Dconditional_simplifier.cc107 HloComputation* new_computation = in TryRemoveUnusedConditionalOperands() local
109 param = new_computation->parameter_instruction(0); in TryRemoveUnusedConditionalOperands()
128 conditional->set_branch_computation(branch, new_computation); in TryRemoveUnusedConditionalOperands()
H A Ddynamic_dimension_inference.cc1483 HloComputation* new_computation = branch_computation; in HandleConditional() local
1490 new_computation, in HandleConditional()
1513 new_computation, dynamic_parameter_binding, parent_)); in HandleConditional()
1515 new_branch_computations.push_back(new_computation); in HandleConditional()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dvalue_inference.cc114 std::unique_ptr<HloComputation> new_computation) { in WithComputation()
115 computation = new_computation.get(); in WithComputation()
120 module.AddEmbeddedComputation(std::move(new_computation)); in WithComputation()
H A Dxla_builder.cc3962 HloComputationProto new_computation(e); in AddCalledComputation() local
3964 remapped_ids[new_computation.id()] = computation_id; in AddCalledComputation()
3965 SetProtoIdAndName(&new_computation, in AddCalledComputation()
3966 GetBaseName(new_computation.name(), kNameSeparator), in AddCalledComputation()
3968 for (auto& instruction : *new_computation.mutable_instructions()) { in AddCalledComputation()
3975 new_computation.set_root_id(remapped_ids.at(new_computation.root_id())); in AddCalledComputation()
3977 imported_computations.push_back(std::move(new_computation)); in AddCalledComputation()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner.cc4201 auto new_computation = in DoPartition() local
4204 DoCodeMotionForWindowedDotGeneralLoops(new_computation, options)); in DoPartition()
4208 replacement[computation] = new_computation; in DoPartition()