llvm::StringMap<PartitionedCallStackOpsInfo>*);

// Handles stack usage by a tf.While. It will convert the body and conditional
// function signatures, and performs stack ops decomposition on them.
LogicalResult HandleWhileOp(
    TF::WhileOp while_op, ModuleOp module,
    const llvm::SmallDenseMap<Value, Value>& data_var_to_size_var,
    llvm::StringMap<PartitionedCallStackOpsInfo>*
        decomposed_partitioned_call_callees) {

  auto output_buffer_to_size = AddTensorListSizesToTerminator<func::ReturnOp>(
      func.front(), buffer_to_size);
  UpdateFuncType(func);
  return output_buffer_to_size;
}

LogicalResult HandleWhileOp(
    TF::WhileOp while_op, ModuleOp module,
    llvm::SmallDenseMap<Value, SizeInfo>* buffer_to_size,
    llvm::StringMap<PartitionedCallDecompositionInfo>*
        decomposed_partitioned_call_callees) {