bool NotFromDequant(mlir::Value value) {
  auto dequant_op = value.getDefiningOp<DequantizeOp>();
  if (dequant_op) {
    return false;
  }
  auto split_op = value.getDefiningOp<SplitOp>();
  if (!split_op) {
    return true;
  }
  return !split_op.getValue().getDefiningOp<DequantizeOp>();
}

// Optimize TFLite operations in functions.
class OptimizeBatchMatmulPass

  auto iota_broadcast =
      dyn_cast_or_null<mhlo::BroadcastInDimOp>(iota.getDefiningOp());
  if (!iota_broadcast || iota_broadcast.getBroadcastDimensions() != dimensions)
    return false;
  if (!isa_and_nonnull<mhlo::IotaOp>(
          iota_broadcast.getOperand().getDefiningOp()))
    return false;
  return true;
}

                                            PatternRewriter &rewriter) {
  auto dot = dyn_cast_or_null<mhlo::DotGeneralOp>(
      reshape_after.getOperand().getDefiningOp());
  if (!dot) return failure();

  auto reshape_before =
      dyn_cast_or_null<mhlo::ReshapeOp>(dot.getLhs().getDefiningOp());
  if (!reshape_before) return failure();

  if (!dot.getLhs().getType().hasStaticShape() ||
      !dot.getRhs().getType().hasStaticShape() ||

  Operation *op_input = args[0].getDefiningOp();
  Operation *op_weight = args[1].getDefiningOp();
  if (isa<TF::SubOp>(op_input)) {
    op_input = op_input->getOperand(0).getDefiningOp();
  }
  if (isa<TF::SubOp>(op_weight)) {
    op_weight = op_weight->getOperand(0).getDefiningOp();
  }
  if (isa<TF::CastOp>(op_input)) {
    op_input = op_input->getOperand(0).getDefiningOp();

  LogicalResult matchAndRewrite(TF::LookupTableFindV2Op find_op,
                                PatternRewriter& rewriter) const override {
    auto handle_op = find_op.getTableHandle().getDefiningOp();
    if (handle_op == nullptr) return failure();
    auto hashtable_op = llvm::dyn_cast<TFL::HashtableOp>(handle_op);
    if (hashtable_op == nullptr) return failure();
    rewriter.replaceOpWithNewOp<TFL::HashtableFindOp>(

bool IsQintValueDefinedByIntToQintCast(Value v) {
  if (!IsIllegalType(v.getType())) {
    return true;
  }
  if (!v.getDefiningOp() || !llvm::isa<TF::CastOp>(v.getDefiningOp())) {
    return false;
  }
  return IsIntToQintCast(v.getDefiningOp());
}

bool IsTFUniformQuantizedOpLegal(Operation *op) {
  // UniformQuantized Ops are considered legal if its qint operands and

      if (IsTPUOp(input_value.get().getDefiningOp()) &&
          !backward_pass_ops.contains(input_value.get().getDefiningOp()))
        values_to_add_nodes.insert(input_value.get());
  }

  for (Value value : values_to_add_nodes) {
    builder.setInsertionPointAfter(value.getDefiningOp());
    std::vector<Type> types(num_replicas, value.getType());
    Location loc = value.getDefiningOp()->getLoc();

      resource_alias_analysis.GetResourceAliases(resource_iterator);

  auto is_generator = [](Value val) {
    if (mlir::isa<BlockArgument>(val)) return true;
    Operation* definition = val.getDefiningOp();
    return definition->getNumOperands() == 0 &&
           definition->getNumResults() == 1;
  };

  // Check all generator aliases (ops or function argument) are on CPU.

    // * the arg to the identity op is produced by a TPUExecuteOp.
    if (op_device && op_device.getValue().contains("TPU")) {
      return true;
    }

    Operation *def_op = arg.getDefiningOp();
    // If the arg to this identity op is the arg of a function, there's no
    // defining op.
    if (def_op != nullptr &&
        (isa<TF::TPUExecuteOp, TF::TPUExecuteAndUpdateVariablesOp>(def_op))) {

  output_types.reserve(num_return_operands);
  for (int i = 0; i != num_return_operands; ++i) {
    auto returned_value = terminator->getOperand(i);
    Operation* returned_op = returned_value.getDefiningOp();
    if (returned_op && returned_op->hasOneUse() &&
        llvm::isa<DequantizeOp>(returned_op)) {
      auto dequantize_op = llvm::cast<DequantizeOp>(returned_op);