Refactor to insert explicit reshards for func return and data flow ops on their own methods.

Refactor to insert explicit reshards for func return and data flow ops on their own methods. #1266