module {
// CHECK-LABEL: fuse_map_and_batch
func.func @fuse_map_and_batch() -> tensor<!tf_type.variant> attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "BatchDatasetV2"}} {
  %0 = "tf.Const"() {value = dense<5> : tensor<i64>} : () -> tensor<i64>
  %1 = "tf.Const"() {value = dense<false> : tensor<i1>} : () -> tensor<i1>

    key: "value"
    value {
      tensor {
        dtype: DT_BOOL
        tensor_shape {
        }
        bool_val: false
      }
    }
  }
}
node {
  name: "BatchDatasetV2"
  op: "BatchDatasetV2"
  input: "MapDataset"
  input: "batch_size"
  input: "drop_remainder"
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
          dim {

module {
// CHECK-LABEL: fuse_pmap_and_batch
func.func @fuse_pmap_and_batch() -> tensor<!tf_type.variant> attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "BatchDatasetV2"}} {
  %0 = "tf.Const"() {value = dense<5> : tensor<i64>} : () -> tensor<i64>
  %1 = "tf.Const"() {value = dense<false> : tensor<i1>} : () -> tensor<i1>

  %1 = func.call @id(%arg0) : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
  "tf.MakeIterator"(%1, %0) {_class = ["loc:@BatchDatasetV2"], device = ""} : (tensor<!tf_type.variant>, tensor<!tf_type.resource>) -> ()
  %2:2 = "tf.IteratorGetNext"(%0) {_class = ["loc:@iterator"], device = ""} : (tensor<!tf_type.resource>) -> (tensor<200x28x28x1xf32>, tensor<200x10xf32>)

    Res<TF_FloatTensor, [{4-D.  Gradients w.r.t. the input of `avg_pool`.}]>:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<1>;
}

def TF_BatchDatasetV2Op : TF_Op<"BatchDatasetV2", [Pure]> {
  let summary = [{
Creates a dataset that batches `batch_size` elements from `input_dataset`.
  }];

  let arguments = (ins
    TF_VariantTensor:$input_dataset,