}

// Transform pad(transpose(x)) to transpose(pad(x))
struct TransposeCommuteWithPad : public OpRewritePattern<stablehlo::PadOp> {
  using OpRewritePattern<stablehlo::PadOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(stablehlo::PadOp pad_op,
                                PatternRewriter& rewriter) const override {
    Value pad_input = pad_op.getOperand();
    RankedTensorType pad_type = pad_op.getType().cast<RankedTensorType>();

class XlaCallModuleOpToCallOp : public OpRewritePattern<TF::XlaCallModuleOp> {
 public:
  explicit XlaCallModuleOpToCallOp(MLIRContext* context)
      : OpRewritePattern<TF::XlaCallModuleOp>(context) {}

  LogicalResult matchAndRewrite(TF::XlaCallModuleOp op,
                                PatternRewriter& rewriter) const override {
    auto module_op = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module_op);

class TFQuantTypePattern : public ConversionPattern {
 public:
  TFQuantTypePattern(MLIRContext *ctx, TypeConverter &converter)
      : ConversionPattern(converter, MatchAnyOpTypeTag(), 1, ctx) {}

  LogicalResult matchAndRewrite(
      Operation *op, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    // This pattern only handle non-UQ, non-const ops.

 public:
  explicit LowerAddNOp(MLIRContext *context)
      : RewritePattern(AddNOp::getOperationName(), 1, context,
                       {AddV2Op::getOperationName()}) {}

  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const override {
    auto addn_op = cast<AddNOp>(op);

    // TODO(hinsu): Support variant with TensorList type. tf.AddV2 doesn't

template <typename SourceOp>
class PrepareLstmOutputScale : public OpRewritePattern<SourceOp> {
 public:
  explicit PrepareLstmOutputScale(MLIRContext* context)
      : OpRewritePattern<SourceOp>(context) {}
  LogicalResult matchAndRewrite(SourceOp op,
                                PatternRewriter& rewriter) const override {
    operator_property::OpVariant lstm_variant;
    operator_property::OperatorProperty lstm_property;

}

class UnfuseBatchNormTrainingPattern
    : public OpRewritePattern<mhlo::BatchNormTrainingOp> {
 public:
  using OpRewritePattern<mhlo::BatchNormTrainingOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(mhlo::BatchNormTrainingOp bn_op,
                                PatternRewriter &rewriter) const override {
    auto inputs = bn_op.getOperand();
    auto input_type = mlir::dyn_cast<RankedTensorType>(inputs.getType());

 public:
  BFloat16TypePattern(TypeConverter& converter, MLIRContext* ctx)
      : ConversionPattern(converter, MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}

  LogicalResult matchAndRewrite(
      Operation* op, const ArrayRef<Value> operands,
      ConversionPatternRewriter& rewriter) const override {
    if (getTypeConverter()->isLegal(op)) {
      return failure();
    }

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