mhlo::ReturnOp return_op = dyn_cast<mhlo::ReturnOp>(body.back());
  if (!return_op || return_op.getNumOperands() != 2) return failure();

  mhlo::SelectOp value_select = llvm::dyn_cast_or_null<mhlo::SelectOp>(
      return_op.getOperand(0).getDefiningOp());
  if (!value_select || value_select.getOnTrue() != body.getArgument(0) ||
      value_select.getOnFalse() != body.getArgument(2))
    return failure();

               "[](bool v){ return v == " #val# ";}))">>;

// Remove select operators when the result is known in advance.
foreach SelectOp = [TFL_SelectOp, TFL_SelectV2Op] in {
  // select(true_tensor, A, B) -> A
  def Optimize#SelectOp#True : Pat<
    (SelectOp:$result (Arith_ConstantOp $constant),
                      $input1,
                      $input2),
    (replaceWithValue $input1),

          TF::ExpandDimsOp,
          TF::IdentityNOp,
          TF::IdentityOp,
          TF::MaxPoolOp,
          TF::PadV2Op,
          TF::RankOp,
          TF::ReshapeOp,
          TF::SelectOp,
          TF::SelectV2Op,
          TF::ShapeNOp,
          TF::ShapeOp,
          TF::SizeOp,
          TF::SqueezeOp,
          TF::TransposeOp
          // go/keep-sorted end

                mlir::stablehlo::DynamicSliceOp, mlir::stablehlo::GatherOp,
                mlir::stablehlo::PadOp, mlir::stablehlo::ReduceWindowOp,
                mlir::stablehlo::ReshapeOp, mlir::stablehlo::SelectOp,
                mlir::stablehlo::SliceOp, mlir::stablehlo::TransposeOp>(op)) {
    scale_spec->has_same_scale_requirement = true;
  }
  if (llvm::isa<mlir::stablehlo::DynamicSliceOp, mlir::stablehlo::GatherOp,

def CreateTFReadVariableOpFromResourceHandle : NativeCodeCall<
    "$_builder.create<TF::ReadVariableOp>("
    "$0.getLoc(), GetResourceSubtype($1), $1)">;

def CreateTFSelectOp: NativeCodeCall<
    "$_builder.create<TF::SelectOp>("
    "$0.getLoc(), $3.getType(), $1, $2, $3)">;

def ConstAttrIfThenElse: NativeCodeCall<
    "($0).getValue() ? ($1) : ($2)">;

// Convert clamp(lo, x, hi) to clipbyvalue(x, lo, hi).
def Clamp: NativeCodeCall<

  Value selected_input =
      b.create<SelectOp>(input_type, compare_dt, lhs_val, rhs_val);

  Value compare_eq =
      b.create<CompareOp>(lhs_val, rhs_val, ComparisonDirection::EQ);
  Value min_index = b.create<MinOp>(lhs_index, rhs_index);
  Value min_val_index =
      b.create<SelectOp>(index_type, compare_dt, lhs_index, rhs_index);
  Value selected_index =
      b.create<SelectOp>(index_type, compare_eq, min_index, min_val_index);

// TODO: b/322428814 - Add StableHLO quantizer integration tests for ODML.
class RewriteQuantizedSelectOp : public OpRewritePattern<stablehlo::SelectOp> {
 public:
  using OpRewritePattern<stablehlo::SelectOp>::OpRewritePattern;

  LogicalResult match(stablehlo::SelectOp op) const override {
    if (!IsQuantizedTensorType(op.getOperand(1).getType())) {
      return failure();
    }

      TypeID::get<TF::ReciprocalOp>(),
      TypeID::get<TF::ReluOp>(),
      TypeID::get<TF::Relu6Op>(),
      TypeID::get<TF::ReluGradOp>(),
      TypeID::get<TF::RsqrtOp>(),
      TypeID::get<TF::SelectOp>(),
      TypeID::get<TF::SigmoidOp>(),
      TypeID::get<TF::SignOp>(),
      TypeID::get<TF::SoftmaxOp>(),
      TypeID::get<TF::SqrtOp>(),
      TypeID::get<TF::TanhOp>(),
      TypeID::get<TF::XlaConvV2Op>(),

    return tensor;
  }

  return {};
}

//===----------------------------------------------------------------------===//
// SelectOp
//===----------------------------------------------------------------------===//

// Verifies a few extra requirements on SelectOp:
// (1) `then` and `else` must have same shape
// (2) At least one of the following must be true: