if (methods_.summarize == nullptr) {
      return tensorflow::CustomDeviceTensorHandle::SummarizeValue(summary);
    }
    TF_Status c_status;
    std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> summary_buffer(
        methods_.summarize(data_, &c_status), TF_DeleteBuffer);
    if (!c_status.status.ok()) {
      return c_status.status;
    }

  // Replace IfOp with PartitionedCallOp or StatefulPartitionedCallOp.
  auto rewrite = [&](auto op_type) {
    auto empty = rewriter.getStringAttr("");
    ReplaceTfOpWithNewOp<typename decltype(op_type)::CallOp>(
        rewriter, op, op.getResultTypes(), op.getInput(), func,
        /*config=*/empty, /*config_proto=*/empty, /*executor_type=*/empty);
  };

  if (op.getIsStateless())