Merge pull request #854 from google/moe_writing_branch_from_3f02dd5fe0748fb8ab9b3f3682b96b6c8e6d7815

Merge internal changes