auto value = device_attr.getValue();
  // TODO(b/229028654): Remove string conversion once we have C++17.
  absl::string_view device_type(value.data(), value.size());
  // Device type may be empty for some ops, e.g. tf.PartitionedCall.
  auto it = std::find(kValidDeviceTypes.begin(), kValidDeviceTypes.end(),
                      device_type);
  if (it == kValidDeviceTypes.end()) return failure();
  return success();
}

  // CHECK-NOT: Optional
  %0 = "tf.OptionalFromValue"(%arg0) : (tensor<f32>) -> tensor<!tf_type.variant<tensor<f32>>>
  %1 = "tf.PartitionedCall"(%0) {
      config = "",
      config_proto = "",
      executor_type = "",
      f = @identity} : (tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>>

    %0 = tf_executor.graph {
      %out, %ctl = tf_executor.island wraps "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @serving_default} : (tensor<?x!tf_type.string>) -> tensor<*xi64>
      tf_executor.fetch %out : tensor<*xi64>
    }
    return %0 : tensor<*xi64>
  }

      return failure();
    }
    rewriter.setInsertionPointAfter(call_op);

    // The TPUPartitionedCall has a TPUOrdinalSelectorOp for its last argument
    // which should be removed. So the replaced PartitionedCall op should keep
    // its original arguments except for the last element.
    SmallVector<Value> args = call_op.getOperands().drop_back();

    rewriter.replaceOpWithNewOp<TF::PartitionedCallOp>(

    .SetShapeFn([](InferenceContext* c) {
      c->set_output(0, c->input(0));
      return absl::OkStatus();
    })
    .Doc(R"(XLA Merge Op. For use by the XLA JIT only.

Merges the outputs from the PartitionedCall node and the _XlaRun node.
Unlike the TensorFlow Merge op, which requires inputs of some types to be
placed on the host, the _XlaMerge op can merge inputs of all types when

        tf_executor.graph {
          %control_0 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_1}> : () -> ()
          %control_1 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_2}> : () -> ()

  attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    %0 = "tf.PartitionedCall"(%arg0) {f = @__inference_main} : (tensor<1x5x5x1024xf32>) -> tensor<1x5x5x1024xf32>
    return %0 : tensor<1x5x5x1024xf32>
  }

// CHECK-NEXT: [[q_bias:%.+]] = "quantfork.qcast"([[bias]]) : (tensor<2xf32>) -> tensor<2x!quant.uniform<i32:f32, 0.044022349891595126>>

      return true;
    }
  }
  return false;
}

// Checks if the function is only used by supported ops. Returns false when the
// function has no uses. Currently, only PartitionedCall is supported.
// TODO(b/284222309): Support lifting for functions called by control flow.
bool UsedBySupportedOps(ModuleOp module, func::FuncOp func) {
  auto function_uses =

  // Collect all the ops that needs to have token input names attributes. These
  // ops are communication ops and all their parent ops via nesting or function
  // calls. For example, IfRegion op and PartitionedCall op.
  std::vector<Operation*> worklist;
  absl::flat_hash_set<Operation*> ops_with_tokens;
  module.walk([&](Operation* op) {
    if (IsCommunicationOp(op)) {
      ops_with_tokens.insert(op);