Type ref_qtype = nullptr;
  for (auto *var_handle_user : var_handle_op.getResult().getUsers()) {
    auto read_variable_op = dyn_cast_or_null<ReadVariableOp>(var_handle_user);
    if (!read_variable_op) continue;
    for (auto *read_variable_user : read_variable_op.getResult().getUsers()) {
      auto q_op = dyn_cast_or_null<QuantizeOp>(read_variable_user);
      if (!q_op || ref_qtype) continue;

    if (!function_name.starts_with(kDequantizeFunctionName)) return failure();

    llvm::SmallVector<Operation*> users(op->getUsers().begin(),
                                        op->getUsers().end());

    bool changed = false;
    for (auto& use : op->getUses()) {
      Operation* user_op = use.getOwner();
      int user_idx = use.getOperandNumber();
      if (!IsOpWithInt8TypeOperand(user_op)) continue;

  if (!result.hasOneUse()) return resource;
  Operation* result_user = *result.getUsers().begin();
  auto assign_var = dyn_cast<TF::AssignVariableOp>(result_user);
  if (!assign_var) return resource;

  auto handle = assign_var.getResource();
  // Skip result if cluster writes to the same variable via multiple results.
  for (Operation* handle_user : handle.getUsers()) {
    if (handle_user == assign_var) continue;

    resource_users.emplace_back(op);
  };

  llvm::SmallVector<Operation*, 4> resource_users_to_visit;
  for (auto composite_arguments : replicate.GetPackedBlockArguments()) {
    for (auto resource_user : composite_arguments.getUsers())
      resource_users_to_visit.emplace_back(resource_user);
  }

  while (!resource_users_to_visit.empty()) {
    llvm::SmallVector<Operation*, 4> new_resource_users;

  if (!classes_attr) {
    // Attempt to parse "_class" from the IdentityOp that follows VariableV2.
    // For read-only reference variables, IdentityOp should be the only user of
    // VariableV2.
    auto identity_op = op->getUsers().begin();
    classes_attr = identity_op->getAttrOfType<ArrayAttr>(kClassAttr);
    if (!classes_attr) {
      op->emitOpError() << "has no '_class' and 'shared_name' attributes";
      return StringRef();
    }
  }

        freezeable[val] = false;
        continue;
      }

      freezeable[val] = true;

      // Verify users are supported kind.
      for (Operation *user : val.getUsers()) {
        if (!(isa<TF::ReadVariableOp>(user) || isa<CallOpInterface>(user))) {
          freezeable[val] = false;
          // Error out early if possible.
          if (!allow_mutable_tensors) {

    // future.
    OpBuilder b(op);
    RankedTensorType type = RankedTensorType::get({}, b.getIntegerType(32));
    int index = op.getDeviceNames().size();
    for (auto use : op.getOperation()->getUsers()) {
      // Skip if it doesn't feed into case. Alternatively this could always
      // return the CPU device index if it exists.
      if (!isa<TF::CaseOp>(use)) return;
    }
    DenseElementsAttr attr =

namespace mlir {
namespace quant {

LogicalResult RemoveIdentity::matchAndRewrite(TF::IdentityOp identity,
                                              PatternRewriter &rewriter) const {
  for (Operation *user : identity->getUsers()) {
    // Replace the op with the input if output is only used by TF ops.
    // Currently this is more on the conservative side since we need to ensure
    // every consumer op to be a TF op before applying this pattern. We can

  // a DenseHostTensor. So as long as one of the users of a captured variable is
  // a fallback op, we should keep this variable as a fallback tensor.
  for (mlir::Operation *user : arg.getUsers()) {
    if (llvm::isa<FallbackAsyncDialect>(user->getDialect())) return true;
  }
  return false;
}

void ForEachArgConsumedByFallback(
    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action) {

      return mlir::failure();
    }

    for (auto result : old_parallel_execute.getExecuteOutputs()) {
      for (mlir::Operation* user :
           llvm::make_early_inc_range(result.getUsers())) {
        if (llvm::isa<mlir::TF::TPUPartitionedOutputV2Op>(user)) {
          assert(user->use_empty());
          user->erase();
        }
      }
    }

    for (auto operand : cluster.getOperands()) {