new_shape, &rewriter);
    reshape_ops.push_back(reshape_op.getResult());
  }

  // Deal with the axis.
  // We don't need to handle axis < 0, since it's counting reversely.
  int32_t axis = concat_op.getAxis();
  if (axis >= 0) {
    axis += (4 - rank);
  }

  // Replace with the new concat op.
  SmallVector<int64_t, 4> new_output_shape;
  for (int i = 0; i < 4 - rank; ++i) {

    %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, tac.device = "CPU", values_count = 2 : i32} : (tensor<100xf32>, tensor<100xf32>) -> tensor<2x100xf32>
    func.return %0 : tensor<2x100xf32>
  }

    The new axis is created at dimension `axis` (default: the new axis is
    appended at the end).
  }];

  let arguments = (ins
    TFL_TensorOf<[I32, I64]>:$indices,
    TFL_I32Tensor:$depth,
    TFL_TensorOf<[F32, I32, I64, I1, I8, UI8]>:$on_value,
    TFL_TensorOf<[F32, I32, I64, I1, I8, UI8]>:$off_value,

    I32Attr:$axis
  );

  let results = (outs

}

// Returns true iff `type` is a uniform quantized type whose storage type is
// 8-bit integer and expressed type is f32.
bool IsI8F32UniformQuantizedType(Type type);

// Returns true iff `type` is a uniform quantized per-axis (per-channel) type
// whose storage type is 8-bit integer and expressed type is f32.
bool IsI8F32UniformQuantizedPerAxisType(Type type);

// Returns true iff `type` is a uniform quantized type whose storage type is

  int64_t value_rank = value_type.getRank();
  int64_t axis = op.getAxis();
  if (axis < -value_rank || axis >= value_rank)
    return op.emitOpError("axis attribute must be in the range of [-")
           << value_rank << ", " << value_rank << ')';

  axis = GetDimForAxis(axis, value_rank);
  int64_t dim_size = value_type.getDimSize(axis);
  if (ShapedType::isDynamic(dim_size)) return success();

        spec->biases_params[2] = {{0, 1},
                                  quant::GetUniformQuantizedTypeForBias};
      }
    } else if (function_name.contains("gather")) {
      // Note that gather has axis attribute that specifies channel axis.
      spec->coeff_op_quant_dim[0] = -1;
    }
    for (auto quantizable_operand : spec->coeff_op_quant_dim) {
      spec->quantizable_operands.insert(quantizable_operand.first);
    }
  }

  Type packed_type = RankedTensorType::get(
      {bcast.output_batch_size(), rows, cols}, element_type);
  const auto axis = rewriter.getI64IntegerAttr(0);
  auto pack_op =
      rewriter.create<TF::PackOp>(loc, packed_type, /*values=*/matmuls, axis);

  // Reshape the rank-3 tensor into the correct output shape.
  const auto& result_batch_shape = bcast.output_batch_shape().dim_sizes();

              {static_cast<int>(
                  dot_dimensions_info.batch_dimensions().AxesArray().size())},
              builder.getIntegerType(32)),
          operand_shape, batch_axes_tensor, /*axis*/ 0, /*batch_dims*/ 0);
      flattend_shape_values.push_back(batch_dims);
    } else {
      llvm::SmallVector<int32_t> batch_i32_vec;
      for (int64_t element :

// -----

// CHECK-LABEL: @RemoveRedundantUnpackPack
func.func @RemoveRedundantUnpackPack(%arg0: tensor<2x5xf32>) -> tensor<2x5xf32> {
  %0:2 = "tfl.unpack"(%arg0) {axis = 0 : i32, num = 2 : i32} : (tensor<2x5xf32>) -> (tensor<5xf32>, tensor<5xf32>)
  %1 = "tfl.pack"(%0#0, %0#1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<5xf32>, tensor<5xf32>) -> (tensor<2x5xf32>)
  func.return %1: tensor<2x5xf32>
  // CHECK-NOT: pack

  func.return
}

// CHECK-LABEL: testGatherToV2
// Ensures that axis param and batch_dims attr use their default values of 0.
func.func @testGatherToV2(%params: tensor<4x3xf32>, %indices: tensor<1x2xi32>) -> tensor<2x3xf32> {
  // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>