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;

      [&](TF::AnonymousIteratorOp op) { count = replace(op, count); });
  getOperation().walk(
      [&](TF::AnonymousIteratorV2Op op) { count = replace(op, count); });
  getOperation().walk(
      [&](TF::AnonymousIteratorV3Op op) { count = replace(op, count); });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateNameAnonymousIteratorsPass() {
  return std::make_unique<NameAnonymousIteratorsPass>();
}