SmallVector<int64_t, 3> slice_size = {1, num_rows, num_cols};
    Type slice_result_type = RankedTensorType::get(slice_size, element_type);
    llvm::SmallVector<Type, 4> output_types(batch_size, slice_result_type);
    auto split_op = rewriter.create<TF::SplitOp>(loc, output_types,
                                                 split_dimension_op.getOutput(),
                                                 reshape_op.getOutput());

              "max_percentile",
              rewriter.getF32FloatAttr(
                  calib_opts_.calibration_parameters().max_percentile())),
      };

      SmallVector<Type, 4> output_types{
          value.getType(),
          RankedTensorType::get({}, rewriter.getF32Type()),
          RankedTensorType::get({}, rewriter.getF32Type()),

  %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = ""} : () -> tensor<i32>
  %0 = "tf.MultiDeviceIteratorFromStringHandle"(%arg0) <{output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]}> {device = ""} : (tensor<!tf_type.string>) -> tensor<!tf_type.resource>

                            bool include_attr) const;

  string op_name;
  std::vector<string> arg_types;
  std::vector<string> arg_names;
  std::vector<string> output_types;
  std::vector<string> output_names;
  std::vector<bool> is_list_output;
  bool has_optional_attrs;
  string comment;

  const OpDef& graph_op_def;
  const ApiDef& api_def;

      finalize_func = @finalize,
      init_func = @init,
      next_func = @next,
      operandSegmentSizes = array<i32: 1, 2, 1>,
      output_shapes = [#tf_type.shape<>],
      output_types = [!tf_type.string],
      metadata = ""} : (
              tensor<4xf32>,
              tensor<3xf32>,
              tensor<!tf_type.resource>,
              tensor<2xf32>) -> tensor<!tf_type.variant>
  return

  %2 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
  %3 = "tf.RangeDataset"(%0, %1, %2) {__op_key = 3: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
  // CHECK: tf_mlrt.executeop{{.*}}op: \22FlatMapDataset\22
  // CHECK-SAME: \22__inference_Dataset_flat_map_lambda_19\22

  for (Value v : inputs) operand_types.emplace_back(v.getType());
  llvm::SmallVector<Type, 4> output_types;
  output_types.reserve(outputs.size());
  for (Value v : outputs) output_types.emplace_back(v.getType());

  auto func_type = builder->getFunctionType(operand_types, output_types);

  FuncOp outlined_func =
      FuncOp::create(ops.front()->getLoc(), kHostFunctionAttr, func_type);

mlir::LogicalResult GetOutputTypesForLogicalDeviceComputation(
    int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config,
    mlir::tf_device::ClusterFuncOp cluster_func,
    llvm::SmallVectorImpl<mlir::Type>* output_types,
    llvm::SmallVectorImpl<int>* cluster_to_core_index);

// Remaps outputs of `new_parallel_execute` op that represent concurrent
// execution of the `tf_device.cluster_func` at index `cluster_idx` of

  replicas.reserve(num_replicas);
  auto devices = replicate_op.getDevices();

  // Collect result types and operands.
  Operation& terminator = replicate_op.GetBody().back();
  llvm::SmallVector<Type, 8> output_types(terminator.getOperandTypes());
  auto control_type = tf_executor::ControlType::get(island_op.getContext());
  llvm::SmallVector<Value, 8> replica_inputs(island_op.getControlInputs());

      }
      output_types = {new_output_type};
    } else {
      output_types = {output_type.clone(elem_type.getStorageType())};
    }

    SmallVector<Value> args = {q_op.getArg(), scale, zero_point};
    FlatSymbolRefAttr func_name =
        FlatSymbolRefAttr::get(rewriter.getStringAttr(kQuantizeFuncName));

    auto quantize_call = rewriter.create<TF::PartitionedCallOp>(
        loc, output_types, args, func_name,