LogicalResult LowerPackIntoConcatReshape::matchAndRewrite(
    TFL::PackOp pack_op, PatternRewriter& rewriter) const {
  // Pack op should have same shape type.
  SmallVector<Value, 5> pack_inputs(pack_op.getValues());
  auto input_type = mlir::dyn_cast<RankedTensorType>(pack_inputs[0].getType());
  if (!input_type) return failure();

  // Figure out output shapes.
  SmallVector<int64_t, 4> concat_out_shape;
  SmallVector<int64_t, 4> pack_out_shape;

  EXPECT_THAT(pack_input1->quantization->scale[0],
              FloatEq(pack_input2->quantization->scale[0]));
  EXPECT_THAT(pack_input0->quantization->zero_point[0],
              Eq(pack_input1->quantization->zero_point[0]));
  EXPECT_THAT(pack_input1->quantization->zero_point[0],
              Eq(pack_input2->quantization->zero_point[0]));

  EXPECT_THAT(pack_input1->quantization->scale[0],

      replicated_inputs;
  llvm::SmallVector<OpAsmParser::UnresolvedOperand, 8> packed_inputs;
  llvm::SmallVector<OpAsmParser::UnresolvedOperand, 8> region_args;
  llvm::SmallVector<Type, 8> region_arg_types;
  int32_t n = 0;
  Region& body = *result.addRegion();
  if (ParseReplicateOpOperands(&parser, &result, &replicated_inputs,
                               &packed_inputs, &region_args,
                               &region_arg_types) ||

      "llvm::ArrayRef<std::pair<ValueRange, Type>>":$replicated_inputs,
      "ValueRange":$packed_inputs, "TypeRange":$replica_output_types)>,
    OpBuilder<(ins "int":$n, "std::optional<DictionaryAttr>":$devices,
      "llvm::ArrayRef<std::pair<ValueRange, Type>>":$replicated_inputs,
      "ValueRange":$packed_inputs, "TypeRange":$replica_output_types)>,
  ];

  let hasCustomAssemblyFormat = 1;
  let hasVerifier = 1;

  // `num_replicas` or 1. Collect all their operands and associated type for
  // creating the replicate op.
  llvm::SmallVector<std::pair<ValueRange, Type>, 8> replicated_inputs;
  llvm::SmallVector<Value, 8> packed_inputs;
  llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> replicated_ops;
  llvm::SmallVector<mlir::TF::TPUReplicatedInputOp, 8> packed_ops;
  for (const auto& pos_and_input : llvm::enumerate(replicated_input_ops)) {

    for (auto input_output :
         llvm::zip(pack_op.getOperands(), input_unpack_op.getResults())) {
      Value pack_input = std::get<0>(input_output);
      Value unpack_output = std::get<1>(input_output);
      // Make sure the ordering is the same for the pack op & unpack op.
      if (pack_input != unpack_output) return failure();
    }

    // Replace the pack's output to the unpack's input.