Name("GeneratorDataset").Device(DEVICE).HostMemory("handle"),            \
      data::GeneratorDatasetOp);                                               \
  REGISTER_KERNEL_BUILDER(Name("PrefetchDataset")                              \
                              .Device(DEVICE)                                  \
                              .HostMemory("buffer_size")                       \

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;

  let hasFolder = 1;
}

def TF_PrefetchDatasetOp : TF_Op<"PrefetchDataset", [Pure]> {
  let summary = [{
Creates a dataset that asynchronously prefetches elements from `input_dataset`.
  }];

  let arguments = (ins
    TF_VariantTensor:$input_dataset,