op: "TensorSliceDataset"
  input: "tensors/normalize_tensors/component_0"
  attr {
    key: "Toutput_types"
    value {
      list {
        type: DT_INT32
      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
        }
      }
    }
  }
}
node {
  name: "MapDataset"
  op: "MapDataset"
  input: "TensorSliceDataset"
  attr {
    key: "Targuments"

  std::array<int64_t, 4> output_shape;
  // NHWC <- NCHW
  output_shape[0] = composite_result_shape[0];
  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];
  output_shape[3] = composite_result_shape[1];

  auto input_type = mlir::cast<ShapedType>(old_op->getOperand(0).getType());

  return RankedTensorType::get(output_shape, input_type.getElementType());
}

    }
  }
  attr {
    key: "index"
    value {
      i: 4
    }
  }
}
node {
  name: "IteratorGetNext"
  op: "IteratorGetNext"
  input: "args_0"
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
          dim {
            size: 1
          }
          dim {
            size: 8
          }
        }
        shape {
          dim {

      %outputs_4, %control_5 = tf_executor.island wraps "tf.Identity"(%outputs_2) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>

// Transpose conv supports hybrid computation with quantized weights.
def FoldQuantWeightsIntoTposeConv : Pat<
  (TFL_TransposeConvOp
    $output_shape,
    (TFL_DequantizeOp $quant_weights),
    $quant_input,
    $bias, $padding, $stride_h, $stride_w, $faf),
  (TFL_TransposeConvOp $output_shape, $quant_weights,
    $quant_input, $bias, $padding, $stride_h, $stride_w, $faf), 

// cases, the additional attribute should be dropped.

// CHECK: name: "StatefulIf"
// CHECK-NOT: name:
// CHECK: op: "If"
// CHECK-NOT: is_stateless
// CHECK:  attr {
// CHECK:    key: "output_shapes"
// CHECK:    value {
// CHECK:      list {
// CHECK:        shape {
// CHECK:          dim {
// CHECK:            size: 4
// CHECK:          }
// CHECK:        }
// CHECK:      }
// CHECK:    }

  } : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>

  func.return %0 : tensor<!tf_type.variant>
// CHECK: "tf.MapDataset"(
// CHECK-SAME: <{f = @{{.*}}, metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string], preserve_cardinality = true, use_inter_op_parallelism = true}> {Targuments = []}
}

// CHECK-LABEL: TfTakeWhileDataset

def FuseMapAndBatch : Pat<
  (TF_BatchDatasetV2Op
     (TF_MapDatasetOp $input_dataset, $other_arguments, $f, $output_types,
        $output_shapes, $use_inter_op_parallelism, $preserve_cardinality,
        $force_synchronous, $map_dataset_metadata),
     $batch_size, $drop_remainder, $parallel_copy, $batch_output_types,

  attr {
    key: "branches"
    value {
      list {
        func {
          name: "foo"
        }
        func {
          name: "bar"
        }
      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
      }
    }
  }
  experimental_debug_info {
  }
}
library {
  function {
    signature {
      name: "bar"
      output_arg {
        name: "bar"

  %1 = "tf.DeviceIndex"() {device = "", device_names = ["CPU", "GPU"]} : () -> tensor<i32>
  %4 = "tf.Case"(%1, %arg0, %arg1) {branches = [@sub, @add], output_shapes = [#tf_type.shape<>], is_stateless = false} : (tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<f32>

  func.return %0, %4 : tensor<i32>, tensor<f32>
}