!indices_ty.getElementType().isInteger(32))
      return rewriter.notifyMatchFailure(
          op,
          "only match for the case where the second operand an I32 shapeType");
    auto sort_dim = op.getDimension();
    auto k = indices_ty.getDimSize(sort_dim);
    auto rank = keys_ty.getRank();
    if (sort_dim != rank - 1 || k < 1)
      return rewriter.notifyMatchFailure(

        auto concat_option = tflite::CreateStablehloConcatenateOptions(
            builder_,
            mlir::cast<mlir::vhlo::IntegerV1Attr>(vhlo_op.getDimension())
                .getValue()
                .getSExtValue());

        return tflite::CreateOperator(
            builder_, opcode_index, builder_.CreateVector(operands),

  }

  void rewrite(stablehlo::ConcatenateOp op,
               PatternRewriter& rewriter) const override {
    Type output_type = op.getResult().getType();
    uint32_t axis = CastI64ToI32(op.getDimension()).value();
    rewriter.replaceOpWithNewOp<TFL::ConcatenationOp>(
        op, output_type, op.getOperands(), axis,
        /*fused_activation_function=*/rewriter.getStringAttr("NONE"));
  }
};

               << "] different than the shape of input[0]: "
               << input_ty_0.getShape();
      }
    }
  }

  ElementsAttr dimension;
  if (matchPattern(op.getDimension(), m_Constant(&dimension))) {
    if (dimension.getShapedType().getRank() != 0 ||
        dimension.getShapedType().getNumElements() != 1)
      return op.emitOpError() << "dimension must be a scalar";
  }