}];

  let arguments = (ins);

  let results = (outs
    TF_Int32Tensor:$device_ordinal
  );
}

def TF_TakeDatasetOp : TF_Op<"TakeDataset", [Pure]> {
  let summary = [{
Creates a dataset that contains `count` elements from the `input_dataset`.
  }];

  let arguments = (ins
    TF_VariantTensor:$input_dataset,