registry.insert<tf_device::TensorFlowDeviceDialect>();
  }

  void runOnOperation() override;
};

// Create the AnonymousIterator for `reduce_dataset` with `dataset_types` using
// `builder`.
AnonymousIteratorV3Op CreateIterator(OpBuilder builder,
                                     llvm::ArrayRef<Type> dataset_types,
                                     ReduceDatasetOp reduce_dataset) {
  llvm::SmallVector<Attribute, 2> shape_attrs;