// Infers the shape for ReduceDatasetOp and its associated reduce function.
  // Returns whether either op or function type was changed.
  bool InferShapeForReduceDataset(ReduceDatasetOp op, int64_t max_iterations);

  // Infers the shape for TakeWhileDatasetOp and its associated predicate
  // function. Returns whether either op or function type was changed.