op {
  graph_op_name: "TakeDataset"
  in_arg {
    name: "count"
    description: <<END
A scalar representing the number of elements from the `input_dataset`
that should be taken. A value of `-1` indicates that all of `input_dataset`
is taken.
END
  }
  summary: "Creates a dataset that contains `count` elements from the `input_dataset`."
}