value_select.getOperand(0).getDefiningOp());
  if (!value_gt ||
      value_gt.getComparisonDirection() != compare_direction_included ||
      value_gt.getLhs() != body.getArgument(0) ||
      value_gt.getRhs() != body.getArgument(2))
    return failure();

  mhlo::SelectOp index_select = llvm::dyn_cast_or_null<mhlo::SelectOp>(
      return_op.getOperand(1).getDefiningOp());
  if (!index_select) return failure();

  SmallVector<ElementsAttr> inputs;
  for (auto operand : op->getOperands()) {
    auto preceding_const_op = operand.getDefiningOp<TF::ConstOp>();
    if (preceding_const_op) {
      inputs.push_back(preceding_const_op.getValue());
      continue;
    }

    Operation* preceding_op = operand.getDefiningOp();
    int preceding_result_id = -1;
    for (auto preceding_result : preceding_op->getResults()) {

// See the function doc in the header file.
def GetPadOpType : NativeCodeCall<
  "GetPadOpType((*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

// See the function doc in the header file.
def GetAvgPoolOpPadAttr: 
  NativeCodeCall<"GetAvgPoolOpPadAttr($_builder, (*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

// Returns true if the provided padding in the composite op can *not* be 

    Value rhs = mul_op.getRhs();
    conv_op = lhs.getDefiningOp<mhlo::ConvolutionOp>();
    if (conv_op == nullptr) {
      return failure();
    }
    filter = conv_op.getRhs().getDefiningOp<mhlo::ConstantOp>();
    if (filter == nullptr) {
      return failure();
    }
    // Try to match static broadcast or dynamic broadcast.
    bcast_or_const_op = rhs.getDefiningOp();
    bool is_dynamic_broadcast =

    TF::TPUReplicatedInputOp replicated_input) {
  if (!llvm::all_of(replicated_input.getInputs(), [](Value input) {
        return llvm::isa_and_nonnull<TF::TPUPartitionedInputV2Op>(
            input.getDefiningOp());
      }))
    return replicated_input.emitOpError()
           << "expects all inputs from 'tf.TPUPartitionedInputV2' ops";

  const auto metadata_iter =

def NoFallbackAttrNotSet : Constraint<CPred<
  "!$0.getDefiningOp()->hasAttr(kNoFallbackAttr)">>;

class FloatValueEquals<string val> : Constraint<CPred<
  "FloatValueEquals($0, " # val # ")">>;

class RankEquals<string rank> : Constraint<CPred<
  "RankEquals($0, " # rank # ")">>;

def IsFusibleWithBias : Constraint<CPred<
  "IsFusibleWithBiasOp($0.getDefiningOp())">>;

// Folds TF IdentityOp with constant input.

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

// `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 :

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