tf_op->emitError("malformed opdef");
        return failure();
      }
      if (tf_op->hasTrait<mlir::OpTrait::AttrSizedOperandSegments>()) {
        // Add derived "operand_segment_sizes" attr to the created operation.
        // TODO(b/146937733): Don't use <void> here.
        set_segment_sizes_attr(input_ranges, op_def->input_arg(),
                               mlir::OpTrait::AttrSizedOperandSegments<

  fetches.emplace_back(fetch_operand);

  auto builder = OpBuilder::atBlockTerminator(&main_graph_op.GetBody());
  builder.create<FetchOp>(main_graph_op.getLoc(), std::move(fetches));
}

// Creates a new Location for the initializer function. This creates a loc by
// attaching a to the initializer function's type so that it is identifiable.
Location CreateInitOpLoc(MLIRContext* ctx, func::FuncOp init_func_ops) {

  if (!buffer_or.ok()) {
    done(buffer_or.status());
    return;
  }

  xla::PjRtBuffer* pjrt_buffer = (*buffer_or).get();
  if (use_pjrt_tensor_buffer_) {
    // Copy the newly created tensor with PjRtTensorBuffer to output device
    // tensor.
    absl::StatusOr<Tensor> t = MakeTensorFromPjRtBuffer(
        device_tensor->dtype(), device_tensor->shape(), std::move(*buffer_or));
    if (!t.ok()) {

  func_input_arg = block->getArgument(0);
  return quantization_func;
}

// Post-actions after adding quantization logics. Post-actions include
// 1) Adding the created function in the symbol table
// 2) Creating a PartitionedCallOp in the main graph that calls the created
//    function.
TF::PartitionedCallOp FinalizeFunctionRegister(
    PatternRewriter& rewriter, Value input, Value output,

  for (Operation* subgraph_op : llvm::reverse(reverse_subgraph)) {
    // Create a deep copy of the subgraph ops' operands to the func op.
    stablehlo_func_op.getBody().begin()->push_back(subgraph_op->clone(mapper));
  }

  SmallVector<Value> result_values;
  for (const Value original_output_value : outputs) {
    // Use the mapped values in the newly created function that correspond to
    // outputs in the original function.

  // Guarantee all functions have one use, which enables more exact shape
  // inference.
  pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
  // Run shape inference so that tf_executor/tf_device ops created later will
  // likely to inherit more concrete types.
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
  pm.addNestedPass<FuncOp>(mlir::TFTPU::CreateTPUPartitionedOpConversionPass());
  pm.addNestedPass<FuncOp>(

      // Create new ConstantOp-ConvertOp-Operation sequences. At this moment,
      // old ConstantOp is guaranteed to have one F32->F16 convert op regardless
      // of its number of users.
      rewriter.setInsertionPointAfter(op);
      // create new F16 constant op in that location
      ConstantOp new_const = rewriter.create<ConstantOp>(

  auto builder = OpBuilder::atBlockTerminator(&main_graph_op.GetBody());
  // Create an IslandOp that will wrap the IdentityOp. Add a control dependency
  // for the newly copied save function.
  auto wrapper_island_op = builder.create<IslandOp>(
      name_loc, TypeRange{main_file_prefix_arg.getType()},
      tf_executor::ControlType::get(&ctx), ValueRange(control_inputs));

  }
  void runOnOperation() override;
};

// Creates a call to function `func` in region `caller_region`. Use `args` as
// the call arguments, and terminate the region with a yield. The arguments are
// cast to the required type before the call. `use_region_args` control whether
// the input arguments are used as is (for IfOp) or block arguments of the same
// type as the input arguments are created and then used as call arguments (for
// While).

    return tokens[0];
  } else {
    return builder.create<AfterAllOp>(loc, original_token.getType(), tokens)
        .getResult();
  }
}

// Replaces `tf._XlaHostComputeMlir` with individual `mhlo.send` and `mhlo.recv`
// ops per operand and result. Unique Channel IDs are assigned per transfer.
// Sink tokens are created across all `mhlo.send` ops first and then by
// all `mhlo.recv` ops.