bool IsOpWithDataMovementTrait(Operation* op) {
  // Supported data movement ops. These ops do not perform any computations and
  // has one result operand.
  return isa<TF::IdentityOp, TF::CastOp, TF::ReshapeOp, TF::XlaShardingOp,
             TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp, TF::ExpandDimsOp,
             TF::SqueezeOp, TF::TransposeOp>(op);
}

bool IsOpWithQuantizableTrait(Operation* op) {

    if (std::get<0>(res_it) != std::get<1>(res_it)) return std::nullopt;

  // There can only be non-truncating cast op's prior to the call.
  for (; it != block.rend(); ++it) {
    CastOp cast = dyn_cast<CastOp>(*it);
    if (!cast || cast.getTruncate()) return std::nullopt;
  }

  return call;
}

using ArgMatcherFn = function_ref<bool(Value, Region&, Value, Region&)>;

    if (isa<TF::VariantType>(elementType)) {
      // We can only replace OptionalGetValue after the inputs have been
      // replaced.
      return failure();
    }
    rewriter.replaceOpWithNewOp<TF::CastOp>(op, op.getResult(0).getType(),
                                            input);
    return success();
  }
};

class HandleOptionalNone : public OpRewritePattern<TF::OptionalNoneOp> {
 public:

      return failure();
    }

    // Cast to float type before the AvgPool op.
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPointAfter(preceding_sc_op);
    auto fcast_op = rewriter.create<TF::CastOp>(
        preceding_sc_op->getLoc(), dq_arg_type.clone(rewriter.getF32Type()),
        preceding_sc_op.getResult());

    // Create a new AvgPool op with float type.

    Value arg = std::get<0>(ArgAndType);
    Type expected_type = std::get<1>(ArgAndType);
    if (arg.getType() != expected_type) {
      arg = builder.create<CastOp>(loc, expected_type, arg,
                                   /*Truncate=*/builder.getBoolAttr(false));
    }
    casted_args.push_back(arg);
  }
  auto call = builder.create<func::CallOp>(loc, func, casted_args);

  // fallback legalization.
  auto cast_value = [&](OpBuilder& builder, Type result_type, ValueRange inputs,
                        Location loc) -> Value {
    return builder.create<mlir::tensor::CastOp>(loc, result_type,
                                                inputs.front());
  };
  addSourceMaterialization(cast_value);
}

void PopulateLegalizeTfWithTf2XlaPatterns(llvm::StringRef device_type,

      Value operand = std::get<0>(operand_and_type);
      Type expected_type = std::get<1>(operand_and_type);
      if (operand.getType() != expected_type) {
        operand = rewriter.create<TF::CastOp>(
            op.getLoc(), expected_type, operand,
            /*Truncate=*/rewriter.getBoolAttr(false));
      }
      casted_operands.push_back(operand);
    }

    auto call = rewriter.create<func::CallOp>(

          next_values_to_visit.push_back(
              while_op.getBody().front().getArgument(operand_number));
          continue;
        }

        if (llvm::isa<mlir::TF::IdentityOp, mlir::TF::CastOp,
                      mlir::TF::ReadVariableOp>(owner)) {
          next_values_to_visit.push_back(use.getOwner()->getResult(0));
          continue;
        }

  while (!worklist.empty()) {
    mlir::Operation* const use = worklist.pop_back_val();

    // Follow the `CastOp`/`IdentityOp`'s users to handle the `RestoreV2` ->
    // (optionally `IdentityOp`) -> `CastOp` `AssignVariableOp` case.
    if (llvm::isa<TF::CastOp, TF::IdentityOp>(use)) {
      llvm::append_range(worklist, use->getUsers());
      continue;
    }

                                       Location conversion_loc) const final {
    if (!mlir::isa<TensorType>(result_type) ||
        !mlir::isa<TensorType>(input.getType()))
      return nullptr;
    return builder.create<TF::CastOp>(conversion_loc, result_type, input,
                                      /*truncate=*/builder.getBoolAttr(false));
  }

  void processInlinedCallBlocks(
      Operation *call,