// 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>>>

      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

// 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 =