public:
  explicit HoistCwiseUnaryOutOfConcat(MLIRContext* context)
      : OpRewritePattern<TF::ConcatV2Op>(context) {}
  LogicalResult matchAndRewrite(TF::ConcatV2Op op,
                                PatternRewriter& rewriter) const override;
};

LogicalResult HoistCwiseUnaryOutOfConcat::matchAndRewrite(
    TF::ConcatV2Op op, PatternRewriter& rewriter) const {

  auto scale_spec = std::make_unique<OpQuantScaleSpec>();
  if (llvm::isa<
          // clang-format off
          // go/keep-sorted start
          TF::AvgPoolOp,
          TF::ConcatOp,
          TF::ConcatV2Op,
          TF::ExpandDimsOp,
          TF::IdentityNOp,
          TF::IdentityOp,
          TF::MaxPoolOp,
          TF::PadV2Op,
          TF::RankOp,
          TF::ReshapeOp,
          TF::SelectOp,

  auto type = tensorflow::GetTypeFromTFTensorShape({length}, dtype);
  auto axis = rewriter.create<ConstOp>(
      loc, GetScalarOfType(rewriter.getIntegerType(64), 0));
  return rewriter.create<ConcatV2Op>(loc, type, ValueRange(vals), axis);
}

// Lowers AddN op to a sequence of AddV2 ops to accumulate operands.
//
// Note that to improve the parallelism, AddN op uses tree-based reduction.

        TypeID::get<TF::CholeskyOp>(),
        TypeID::get<TF::CollectiveReduceV2Op>(),
        TypeID::get<TF::ComplexAbsOp>(),
        TypeID::get<TF::ConjugateTransposeOp>(),
        TypeID::get<TF::ConcatV2Op>(),
        TypeID::get<TF::ConvOp>(),
        TypeID::get<TF::CoshOp>(),
        TypeID::get<TF::CrossOp>(),
        TypeID::get<TF::CumulativeLogsumexpOp>(),
        TypeID::get<TF::DataFormatDimMapOp>(),

          TF::ArgMaxOp::getOperationName().str(),
          TF::AvgPoolOp::getOperationName().str(),
          TF::BiasAddOp::getOperationName().str(),
          TF::BucketizeOp::getOperationName().str(),
          TF::ConcatV2Op::getOperationName().str(),
          TF::Conv2DBackpropInputOp::getOperationName().str(),
          TF::Conv2DOp::getOperationName().str(),
          TF::DepthwiseConv2dNativeOp::getOperationName().str(),

  // Create a concat of index and trailing zeros.
  llvm::SmallVector<int64_t, 8> zeros(buffer_type.getShape().size() - 1, 0);
  auto zeros_tensor = GetR1Const(zeros, builder, loc);
  return builder.create<TF::ConcatV2Op>(
      loc,
      ArrayRef<Type>{tensorflow::GetTypeFromTFTensorShape(
          {static_cast<int64_t>(buffer_type.getShape().size())},
          getElementTypeOrSelf(index.getType()))},

      CompileTimeConstantOperand<TF::CollectivePermuteOp, 1>,
      // $group_size, $group_key
      CompileTimeConstantOperand<TF::CollectiveReduceV2Op, 1, 2>,
      CompileTimeConstantOperand<TF::ConcatV2Op, -1>,  // (variadic) $axis
      // $filter_sizes
      CompileTimeConstantOperand<TF::Conv2DBackpropFilterOp, 1>,
      CompileTimeConstantOperand<TF::Conv2DBackpropInputOp, 0>,  // $input_sizes
      // $filter_sizes

                                             one_ele_type, 0);

    auto expanded_fft_len_type = tensorflow::GetTypeFromTFTensorShape(
        {2}, fft_len_type.getElementType());

    TF::ConcatV2Op expanded_fft_len = rewriter.create<TF::ConcatV2Op>(
        rfft_op.getLoc(), expanded_fft_len_type,
        SmallVector<Value, 2>({one.getResult(), fft_len}), zero->getResult());

    // Insert the rfft_2d.

  return success();
}

LogicalResult ConvertTFConcatV2Op::matchAndRewrite(
    Operation* op, PatternRewriter& rewriter) const {
  auto tf_concat_op = cast<TF::ConcatV2Op>(op);

  auto values = tf_concat_op.getValues();
  auto output_type = tf_concat_op.getOutput().getType();
  // Extract axis attribute from constant axis tensor
  ElementsAttr axis;

  if (!type.hasRank() || type.getRank() != 1) return {};
  if (port.size() != 2) return {};
  assert(port[0] == 0);
  return ValuePort(op.getOperand(port[1]));
}

ValuePort ComputeInputComponentFor(ConcatV2Op op, ArrayRef<unsigned int> port) {
  if (port.size() != 2) return {};
  assert(port[0] == 0);

  int64_t element_idx = port[1];
  for (Value val : op.getValues()) {