auto output_types = fused_lstm_func_.getFunctionType().getResults();
  SmallVector<int64_t, 2> output_shape{1, mlir::ShapedType::kDynamic};
  EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getShape().size(),
            output_shape.size());
  for (int i = 0; i < output_shape.size(); i++) {
    EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getDimSize(i),
              output_shape[i]);

      }

      // Collect all the quantized outputs and replace them by the results of
      // the new quantized op.
      llvm::SmallDenseMap<Value, int> outputs_replaced;
      SmallVector<Type, 4> output_types;
      output_types.reserve(candidate_op->getNumResults());
      for (const auto& enumerated_result :
           llvm::enumerate(candidate_op->getResults())) {
        Value result = enumerated_result.value();

           output_shapes = [#tf_type.shape<>], output_types = [i32],
           preserve_cardinality = false, sloppy = false,
           use_inter_op_parallelism = true,
           metadata = ""} : (tensor<*x!tf_type.variant>) -> tensor<!tf_type.variant>
  %5 = "tf.BatchDatasetV2"(%4, %0, %1) {device = "",
           output_shapes = [#tf_type.shape<>],
           output_types = [i32],
           parallel_copy = false,

// (inclusive) to `end` index region (exclusive) to `output_types` and returns
// the number of types added.
int AppendTypes(llvm::SmallVectorImpl<Type>* output_types,
                tf_device::ParallelExecuteOp parallel_execute, int start,
                int end) {
  const int size_before = output_types->size();
  for (int index = start; index < end; ++index) {

                                  SmallVectorImpl<Type>* output_types) const;

  // Creates the TF op and also the necessary tfr.cast ops to replace the
  // original TFR call op.
  LogicalResult CreateAndReplaceOp(
      PatternRewriter& rewriter, CallOp call_op,
      const SmallVectorImpl<Type>& output_types,
      const SmallVectorImpl<Value>& inputs, const NamedAttrList& attr_list,

    %batch_size: i64 {tfr.name="batch_size"},
    %num_parallel_calls: i64 {tfr.name="num_parallel_calls"},
    %drop_remainder: i1 {tfr.name="drop_remainder"},
    %f: !tfr.attr {tfr.name="func"},
    %output_types: !tfr.attr {tfr.name="output_types"},
    %output_shapes: !tfr.attr {tfr.name="output_shapes"},
    %preserve_cardinality: i1 {tfr.name="preserve_cardinality", tfr.default=false}) -> !tfr.tensor {

                                     const StringRef func_name,
                                     const TypeRange output_types,
                                     const ValueRange args) {
  TF::PartitionedCallOp call_op = builder.create<TF::PartitionedCallOp>(
      location, output_types, args,
      FlatSymbolRefAttr::get(builder.getStringAttr(func_name)),
      /*config=*/"", /*config_proto=*/"", /*executor_type=*/"");

      Value dequantized_result = dequantize_op.getInput();
      output_types.push_back(dequantized_result.getType());
      terminator->setOperand(i, dequantized_result);
      returned_op->erase();
    } else {
      output_types.push_back(returned_value.getType());
    }
  }
  auto new_func_type = builder.getFunctionType(input_types, output_types);
  func.setType(new_func_type);
}

enum RemoveVolatileOpsType {

  } : (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

    }

    // Outputs.
    const int num_outputs = launch->output_types().size();
    absl::flat_hash_set<Node*> control_outputs;
    std::vector<std::vector<std::pair<Node*, int>>> data_outputs(num_outputs);
    const DataTypeVector& output_types(launch->output_types());

    for (const Edge* le : launch->out_edges()) {
      if (le->IsControlEdge()) {