Home
last modified time | relevance | path

Searched defs:while_op (Results 1 – 25 of 57) sorted by relevance

123

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dwhile_loop_analysis_test.cc55 HloInstruction* while_op = module->entry_computation()->root_instruction(); in TEST_F() local
86 HloInstruction* while_op = module->entry_computation()->root_instruction(); in TEST_F() local
119 HloInstruction* while_op = module->entry_computation()->root_instruction(); in TEST_F() local
154 HloInstruction* while_op = module->entry_computation()->root_instruction(); in TEST_F() local
193 HloInstruction* while_op = module->entry_computation()->root_instruction(); in TEST_F() local
232 HloInstruction* while_op = module->entry_computation()->root_instruction(); in TEST_F() local
267 HloInstruction* while_op = module->entry_computation()->root_instruction(); in TEST_F() local
H A Dwhile_loop_simplifier.cc53 HloInstruction* while_op, in RemoveDeadTupleIndices()
226 static StatusOr<bool> TryRemoveDeadWhileParams(HloInstruction* while_op) { in TryRemoveDeadWhileParams()
472 HloInstruction* while_op, const int64_t tuple_index, in TryRemoveRepeatedWhileTupleIndicesHelper()
523 HloInstruction* while_op) { in TryRemoveRepeatedWhileTupleIndices()
628 static StatusOr<bool> TryRemoveConstantParams(HloInstruction* while_op) { in TryRemoveConstantParams()
776 static StatusOr<bool> TryRemoveWhileLoop(HloInstruction* while_op) { in TryRemoveWhileLoop()
861 static StatusOr<bool> TryPropagateConstant(HloInstruction* while_op) { in TryPropagateConstant()
997 static StatusOr<bool> TryFlattenNestedTuples(HloInstruction* while_op) { in TryFlattenNestedTuples()
1127 HloInstruction* while_op, PrimitiveType elem_ty) { in TryMergeInductionVariables()
H A Dwhile_loop_analysis.cc108 const HloInstruction* while_op) { in GetAuxiliaryLoopInductionVars()
252 optional<int64_t> GetLoopInductionVarTupleIdx(const HloInstruction* while_op) { in GetLoopInductionVarTupleIdx()
388 static optional<int64_t> PatternMatchLoopTripCount(HloInstruction* while_op, in PatternMatchLoopTripCount()
482 optional<int64_t> ComputeWhileLoopTripCount(HloInstruction* while_op, in ComputeWhileLoopTripCount()
572 HloInstruction* while_op) { in ComputeWhileLoopTripCountUpperBound()
H A Dwhile_loop_simplifier_test.cc168 auto* while_op = computation->root_instruction(); in TEST_F() local
185 auto* while_op = computation->root_instruction(); in TEST_F() local
201 auto* while_op = computation->root_instruction(); in TEST_F() local
217 auto* while_op = computation->root_instruction(); in TEST_F() local
232 auto* while_op = computation->root_instruction(); in TEST_F() local
H A Dhlo_evaluator_test.cc5004 HloInstruction* while_op = in TEST_F() local
5053 HloInstruction* while_op = in TEST_F() local
5104 HloInstruction* while_op = in TEST_F() local
5155 HloInstruction* while_op = in TEST_F() local
5212 HloInstruction* while_op = in TEST_F() local
5277 HloInstruction* while_op = in TEST_F() local
5364 HloInstruction* while_op = in TEST_F() local
5413 HloInstruction* while_op = in TEST_F() local
5452 HloInstruction* while_op = in TEST_F() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dtf_saved_model_freeze_variables.cc125 } else if (auto while_op = dyn_cast<TF::WhileOp>(user_op)) { in PropagateUsage() local
131 } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(user_op)) { in PropagateUsage() local
279 T GetUpdatedWhileOp(T while_op, const U& argument_types, in GetUpdatedWhileOp()
385 } else if (auto while_op = dyn_cast<TF::WhileOp>(user_op)) { in FreezeVariables() local
389 } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(user_op)) { in FreezeVariables() local
H A Dtpu_variable_runtime_reformatting.cc99 TF::WhileRegionOp while_op, tf_device::ReplicateOp replicate, in AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping()
335 void HandleReplicateOp(TF::WhileRegionOp while_op, in HandleReplicateOp()
477 module.walk([&](TF::WhileRegionOp while_op) { in runOnOperation()
H A Dfunctional_control_flow_to_regions.cc118 LogicalResult ConvertWhileOp(WhileOp while_op) { in ConvertWhileOp()
149 } else if (auto while_op = llvm::dyn_cast<WhileOp>(op)) { in runOnOperation() local
H A Dtensor_list_ops_decomposition.cc150 TF::WhileOp while_op, ModuleOp module, in HandleWhileOp()
268 TF::WhileRegionOp while_op, ModuleOp module, in HandleWhileRegionOp()
878 } else if (auto while_op = llvm::dyn_cast<TF::WhileOp>(&op)) { in DecomposeTensorListOpsInternal() local
914 } else if (auto while_op = llvm::dyn_cast<TF::WhileRegionOp>(&op)) { in DecomposeTensorListOpsInternal() local
H A Dconvert_control_to_data_outputs.cc339 TF::WhileOp RewriteWhileOp(TF::WhileOp while_op, int num_resource_inputs, in RewriteWhileOp()
430 for (TF::WhileOp while_op : while_callers) { in ConvertControlToDataOutputs() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tools/
H A Dhlo_control_flow_flattening_test.cc71 auto while_op = module->entry_computation()->GetInstructionWithName("while"); in TEST_F() local
132 auto while_op = module->entry_computation()->GetInstructionWithName("while"); in TEST_F() local
186 auto while_op = module->entry_computation()->GetInstructionWithName("while"); in TEST_F() local
242 auto while_op = module->entry_computation()->GetInstructionWithName("while"); in TEST_F() local
698 void CheckWhileBound(HloInstruction* while_op, int expected_bound) { in CheckWhileBound()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Ddtensor_allreduce_sum_optimization.cc263 const int index, mlir::TF::WhileRegionOp while_op, mlir::Value while_output, in CheckWhileLoopOptimizationCriteria()
368 mlir::TF::WhileRegionOp while_op, mlir::OpOperand& add_input, in ExtractAllReduceFromWhileOp()
418 mlir::TF::WhileRegionOp while_op, bool* changed) { in OptimizeWhileLoopLazyAllReduce()
477 if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(op)) in CollectOptimizationCandidates() local
H A Dspmd_expander_common.cc365 } else if (auto while_op = in TraceUseToNextTFOp() local
376 } else if (auto while_op = mlir::dyn_cast<mlir::TF::WhileRegionOp>(owner)) { in TraceUseToNextTFOp() local
749 if (auto while_op = mlir::dyn_cast<mlir::TF::WhileRegionOp>(op)) { in next() local
/aosp_15_r20/external/tensorflow/tensorflow/core/transforms/eliminate_passthrough_iter_args/
H A Dpass.cc193 } else if (auto while_op = dyn_cast<WhileRegionOp>(op)) { in runOnOperation() local
196 } else if (auto while_op = dyn_cast<StatelessWhileRegionOp>(op)) { in runOnOperation() local
199 } else if (auto while_op = dyn_cast<StatefulWhileRegionOp>(op)) { in runOnOperation() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/lite/transforms/
H A Dwhile_loop_outline.cc69 bool IsAlreadyOutlined(WhileOp while_op) { in IsAlreadyOutlined()
190 void WhileOutlinePass::OutlineWhile(WhileOp while_op) { in OutlineWhile()
281 [&](mlir::TFL::WhileOp while_op) { OutlineWhile(while_op); }); in runOnOperation()
H A Dlegalize_tf_while.cc60 void RunOnWhile(TF::WhileOp while_op) { in RunOnWhile()
76 func.getBody().walk([](TF::WhileOp while_op) { RunOnWhile(while_op); }); in RunOnFunction()
H A Dreduce_while_operands.cc186 bool ReduceWhileOperands(TFL::WhileOp while_op) { in ReduceWhileOperands()
292 fn.walk([&](TFL::WhileOp while_op) { ReduceWhileOperands(while_op); }); in runOnOperation()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/transforms/
H A Doptimize_tf_control_flow_side_effect.cc34 if (auto while_op = llvm::dyn_cast<mlir::TF::WhileOp>(op)) { in FunctionHasSideEffect() local
98 if (auto while_op = llvm::dyn_cast<mlir::TF::WhileOp>(op)) { in runOnOperation() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dfunctionalize_control_flow_test.cc395 auto while_op = in TEST() local
545 auto while_op = in TEST() local
663 auto while_op = in TEST() local
812 auto while_op = in TEST() local
1168 auto while_op = ops::While(scope.WithOpName("outer/LoopCond"), in CheckOuterNodesFunctionalized() local
1223 auto while_op = in CheckOuterNodesFunctionalized() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
H A Dlegalize_tf_control_flow.cc192 auto while_op = in LowerWhile() local
375 auto while_op = in LowerWhileRegion() local
422 if (auto while_op = dyn_cast<TF::WhileOp>(op)) { in runOnOperation() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/analysis/
H A Dresource_value_typed_analyzer.cc143 if (auto while_op = dyn_cast<TF::WhileOp>(op)) { in AnalyzeRegion() local
150 if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) { in AnalyzeRegion() local
H A Dresource_alias_analysis.cc382 } else if (auto while_op = dyn_cast<WhileOp>(op)) { in ResourceAliasAnalysisInfo() local
490 Operation* while_op, const BacktrackAnalysisInfo& body_info) { in AnalyzeWhileLoop()
/aosp_15_r20/external/tensorflow/tensorflow/python/ops/
H A Dwhile_v2.py632 def _get_graph(while_op, func_attr_name, attr_graph_name): argument
655 def _create_grad_func(ys, xs, grads, cond_graph, body_graph, name, while_op, argument
766 def _resolve_grad_captures(body_graph, body_grad_graph, while_op): argument
/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Drearrange_function_argument_pass_test.cc104 auto while_op = in TEST() local
217 auto while_op = ops::While(s.WithOpName("while"), in TEST() local
/aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/
H A Dlower_while_op.cc65 static Status Run(Node* while_op, const NameAttrList& cond_fn, in Run()
180 LowerWhileHelper::LowerWhileHelper(Node* while_op, const NameAttrList& cond_fn, in LowerWhileHelper()

123