class ConvertXlaAllReduce
    : public CollectiveRewritePattern<TF::XlaAllReduceOp> {
 public:
  using CollectiveRewritePattern::CollectiveRewritePattern;

  LogicalResult matchAndRewrite(TF::XlaAllReduceOp all_reduce,
                                PatternRewriter& rewriter) const override {
    DenseIntElementsAttr replica_groups;
    if (failed(ConvertReplicaGroups(rewriter, all_reduce.getGroupAssignment(),

  }
}

template <int num_spatial_dims>
class ConvertNdConvOp : public OpConversionPattern<mhlo::ConvolutionOp> {
 public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      mhlo::ConvolutionOp conv_op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    SetDefaultConvAttributes(conv_op, rewriter);
    if (!IsSupportedConvOp(conv_op)) {

// constant tensor with rank 2
struct ConvertBatchMatMulOp2FullyConnectedOp
    : public OpRewritePattern<TFL::BatchMatMulOp> {
  using OpRewritePattern<TFL::BatchMatMulOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(TFL::BatchMatMulOp bmm_op,
                                PatternRewriter& rewriter) const override {
    DenseElementsAttr constant;
    if (auto rhs = bmm_op.getY(); !matchPattern(rhs, m_Constant(&constant))) {

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.

  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,

  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"

  // argument is the orchestrator of the sequence of rewrites. The pattern is
  // expected to interact with it to perform any changes to the IR from here.
  mlir::LogicalResult matchAndRewrite(
      mlir::tfd::ConvertTftToDhtOp op,
      mlir::PatternRewriter& rewriter) const override {
    // Look through the inputs of the ConvertTftToDhtOp.
    mlir::Value convert_op_input_0 = op.getOperand(0);

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;