// tests.
  LogicalResult SetupInputOutputTypesIfNull(OpBuilder builder);

  // Modifies the element types of entry block arguments to be user specified
  // and returns  the new argument types.
  LogicalResult ModifyInputNodes(func::FuncOp func,
                                 llvm::SmallVectorImpl<Type>& new_input_types,
                                 OpBuilder builder);