class RemoveCustomCallWithSharding
    : public OpRewritePattern<stablehlo::CustomCallOp> {
  using OpRewritePattern<stablehlo::CustomCallOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(stablehlo::CustomCallOp op,
                                PatternRewriter &rewriter) const override {
    // Removes the custom call with sharding op if the operand type is the
    // same as the result type.

  static std::vector<Value> sliceInput(Value value, int batch_size,
                                       Location loc, PatternRewriter& rewriter);

  LogicalResult matchAndRewrite(BatchMatMulOpType op,
                                PatternRewriter& rewriter) const override;
};

#define GEN_PASS_DEF_UNROLLBATCHMATMULPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

  return success();
}

}  // namespace

template <typename OpTy>
class SmuggleOpPattern : public OpRewritePattern<OpTy> {
 public:
  using OpRewritePattern<OpTy>::OpRewritePattern;

  LogicalResult matchAndRewrite(OpTy op,
                                PatternRewriter& rewriter) const override {
    return SmuggleOp(op, rewriter);
  }
};

class SmuggleDisallowedOpsPass
    : public PassWrapper<SmuggleDisallowedOpsPass,

template <typename BatchMatMulOpType>
class ConvertTFBatchMatMulToEinsumOp
    : public OpRewritePattern<BatchMatMulOpType> {
  using OpRewritePattern<BatchMatMulOpType>::OpRewritePattern;

  LogicalResult matchAndRewrite(BatchMatMulOpType op,
                                PatternRewriter& rewriter) const override {
    Value input_lhs = op.getX();
    Value input_rhs = op.getY();

    // LHS and RHS must be a ranked tensor type

      : ConversionPattern(converter, MatchAnyOpTypeTag(), /*benefit=*/1, ctx),
        device_type_(device_type),
        prefer_tf2xla_(prefer_tf2xla) {}

  LogicalResult matchAndRewrite(
      Operation* op, ArrayRef<Value> operands,
      ConversionPatternRewriter& rewriter) const override {
    // This pattern is a conversion pattern because we want to specify a type

  SmallVector<Value> results;
  if (failed(FoldOperation(builder, op, results))) {
    return op->getResults();
  }
  return results;
}

LogicalResult ConstantFoldQuantizableOperands::matchAndRewrite(
    Operation* op, PatternRewriter& rewriter) const {
  absl::flat_hash_set<int> quantizable_operands = GetQuantizableOperands(op);
  if (quantizable_operands.empty()) return failure();

  bool has_change = false;

 public:
  explicit CommuteBothInputsTransposedWithEwiseOps(MLIRContext *context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const override {
    if (!IsElementwise(op)) {
      return failure();
    }

    const bool any_blargs = llvm::any_of(

  using OpRewritePattern::OpRewritePattern;

  explicit ReplaceCustomCallWithComposite(MLIRContext *context)
      : OpRewritePattern<mlir::stablehlo::CustomCallOp>(context) {}

  LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
                                PatternRewriter &rewriter) const override {
    auto backendConfig =
        mlir::dyn_cast<DictionaryAttr>(op->getAttr("composite.backend_config"));

    : public OpRewritePattern<quantfork::DequantizeCastOp> {
 public:
  explicit RemoveUnusedQdqPattern(MLIRContext* context)
      : OpRewritePattern<quantfork::DequantizeCastOp>(context) {}
  LogicalResult matchAndRewrite(quantfork::DequantizeCastOp dq_op,
                                PatternRewriter& rewriter) const override {
    auto q_op = dq_op.getArg().getDefiningOp<quantfork::QuantizeCastOp>();
    if (!q_op) return failure();

  // outlives this object.
  explicit ConvertCustomAggregationOpToQuantStats(MLIRContext *context)
      : OpRewritePattern<TF::CustomAggregatorOp>(context) {}

  LogicalResult matchAndRewrite(TF::CustomAggregatorOp op,
                                PatternRewriter &rewriter) const override {
    FloatAttr min = mlir::dyn_cast_or_null<FloatAttr>(op->getAttr("min"));