computeInputsOutputs(ops, &inputs, &outputs);

  std::vector<Type> input_types;
  std::vector<Type> output_types;

  input_types.reserve(inputs.size());
  for (Value v : inputs) {
    input_types.push_back(v.getType());
  }
  output_types.reserve(outputs.size());
  for (Value v : outputs) {
    output_types.push_back(v.getType());
  }

  // Create the function.
  MLIRContext* context = ops[0]->getContext();

                                  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,

          }
          dim {
            size: 16
          }
        }
        shape {
          unknown_rank: true
        }
      }
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_BFLOAT16
        type: DT_FLOAT
        type: DT_DOUBLE
      }
    }
  }
}
node {
  name: "IteratorGetNextSync"

           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<i32>) -> tensor<!tf_type.variant>
  %6 = "tf.BatchDatasetV2"(%5, %0, %1) {device = "",
           output_shapes = [#tf_type.shape<>],
           output_types = [i32],
           parallel_copy = false,

tf_device::ParallelExecuteOp BuildParallelExecuteOp(
    tf_device::ClusterFuncOp cluster_func, OpBuilder* builder) {
  const auto output_types = cluster_func.getResultTypes();
  builder->setInsertionPoint(cluster_func);
  auto parallel_execute = builder->create<tf_device::ParallelExecuteOp>(
      cluster_func.getLoc(), 1, output_types);
  cluster_func->remove();
  auto& block = parallel_execute.GetRegionBlockWithIndex(0);

    %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()) {