return failure();
    }
    // If there are redundant casts, hoist output of raw data op originating op.
    if (preceding_cast.getArg().getDefiningOp()) {
      auto redundant_cast = preceding_cast.getArg().getDefiningOp<CastOp>();
      if (!redundant_cast || redundant_cast.getArg().getType() !=
                                 preceding_cast.getOut().getType()) {
        return failure();
      }

  if (block_arg_num.has_value()) return block_arg_num;

  Value next_input = input;
  auto pad_op = dyn_cast_or_null<TF::PadOp>(next_input.getDefiningOp());
  auto cast_op = dyn_cast_or_null<TF::CastOp>(next_input.getDefiningOp());

  while (pad_op || cast_op) {
    if (pad_op) {
      auto block_arg_num = GetBlockArgNum(pad_op.getInput());
      if (block_arg_num.has_value()) return block_arg_num;

      auto& tensor_property = enumerated_inputs.second;

      Value input = op.getOperand(index);

      if (input.getDefiningOp() == nullptr) continue;

      // TODO(b/172517537): make this work with non-PTQ case.
      if (llvm::isa<func::ConstantOp, arith::ConstantOp, TFL::ConstOp>(
              input.getDefiningOp())) {
        // Tensors with derived scale are biases, and handled in propagation.

    if (any_blargs) {
      return failure();
    }

    auto tpose_arg1 = llvm::dyn_cast_or_null<TFL::TransposeOp>(
        op->getOperand(0).getDefiningOp());
    auto tpose_arg2 = llvm::dyn_cast_or_null<TFL::TransposeOp>(
        op->getOperand(1).getDefiningOp());

    if (!tpose_arg1 || !tpose_arg2) {
      return failure();
    }

    auto tpose_arg1_type =

include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mhlo/IR/hlo_ops.td"
include "stablehlo/dialect/ChloOps.td"

def IsDefinedByConvOrDotGeneralOp : Constraint<Or<[
  CPred<"$0.getDefiningOp<mhlo::ConvolutionOp>()">,
  CPred<"$0.getDefiningOp<mhlo::DotGeneralOp>()">]>>;

def IsNull : Constraint<CPred<"!$0">>;

// This pattern optimizes:
//   conv/dot_general + a + b -> conv/dot_general + (a + b)

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;
}

// `tf_quant.composite_function` attribute.
def IsNotInLiftedFunc :
      Constraint<CPred<"!IsInLiftedFunc($0.getDefiningOp())">>;

// Checks if the value is not inside a StableHLO op with region.
def IsNotInStableHloOpRegion :
      Constraint<CPred<"!IsInStableHloOpRegion($0.getDefiningOp())">>;

// Checks if the given einsum op is supported for XlaDotV2 quantization.
def IsEinsumSupportedByXlaDotV2 :

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