batch_norm =
          dyn_cast_or_null<FusedBatchNormV3Op>(add_op.getX().getDefiningOp());
      if (batch_norm) {
        side_input = add_op.getY();
      } else {
        // Didn't get a FusedBatchNorm on the LHS of the AddV2, try the RHS.
        batch_norm =
            dyn_cast_or_null<FusedBatchNormV3Op>(add_op.getY().getDefiningOp());
        if (!batch_norm) return failure();
        side_input = add_op.getX();
      }

class AreIntSplatValuesEqual<string type> : Constraint<CPred<
  "AreSplatValuesEqual<"# type #">($0, $1)">>;

// Returns defining op of this value.
def GetDefiningOp :  NativeCodeCall<"$0.getDefiningOp()">;

// Clones an operation with new operands while keeping attributes.
def CloneOpWithReplacedOperands : NativeCodeCall<
  "CloneOpWithReplacedOperands("

                                            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() ||

      if (!returned_type ||
          !returned_type.getElementType().isa<quant::QuantizedType>()) {
        continue;
      }

      if (auto returned_op =
              returned_value.get().getDefiningOp<TFR::CastOp>()) {
        auto new_type = returned_type.clone(returned_type.getElementType()
                                                .cast<quant::QuantizedType>()

  for (Value operand : op->getOperands()) {
    if (Operation* pred = operand.getDefiningOp()) {
      int result_number = llvm::cast<OpResult>(operand).getResultNumber();
      if (auto pred_island = llvm::dyn_cast<tf_executor::IslandOp>(pred)) {
        Value yield_operand = pred_island.GetYield().getOperand(result_number);
        predecessors.push_back(yield_operand.getDefiningOp());
      }
    }
  }
  return predecessors;
}

  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();

    if (!fetch.getDefiningOp()) {
      // Skip, because there is no op to add control to (eg: function args).
      continue;
    } else if (fetch.getDefiningOp()->getParentOp() == island_op) {
      // Skip, because it is the same island.
      continue;
    } else if (auto other_island_op = llvm::dyn_cast<tf_executor::IslandOp>(
                   fetch.getDefiningOp())) {

    auto read_var = llvm::dyn_cast_or_null<TF::ReadVariableOp>(
        operand.get().getDefiningOp());
    if (!read_var) continue;
    auto partitioned_input =
        llvm::dyn_cast_or_null<TF::TPUPartitionedInputV2Op>(
            read_var.getResource().getDefiningOp());
    if (!partitioned_input || !AllResourceTypesHaveSubtypes(

  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