Merge pull request #34646 from wwwind:fix_pack

PiperOrigin-RevId: 287779568

tensorflow/lite/tools/optimize/BUILD

@@ -241,6 +241,7 @@ tf_cc_test(
         "//tensorflow/lite/tools/optimize:testdata/minimum.bin",
         "//tensorflow/lite/tools/optimize:testdata/mixed.bin",
         "//tensorflow/lite/tools/optimize:testdata/multi_input_add_reshape.bin",
+        "//tensorflow/lite/tools/optimize:testdata/pack.bin",
         "//tensorflow/lite/tools/optimize:testdata/single_avg_pool_min_minus_5_max_plus_5.bin",
         "//tensorflow/lite/tools/optimize:testdata/single_conv_no_bias.bin",
         "//tensorflow/lite/tools/optimize:testdata/single_conv_weights_min_0_max_plus_10.bin",

tensorflow/lite/tools/optimize/quantize_model_test.cc

@@ -1260,6 +1260,68 @@ TEST_F(QuantizeCustomOpTest, VerifyMixedQuantization) {
   }
 }
 
+class QuantizePackTest : public QuantizeModelTest {
+ protected:
+  QuantizePackTest() {
+    input_model_ = ReadModel(internal::kModelPack);
+    readonly_model_ = input_model_->GetModel();
+    readonly_model_->UnPackTo(&model_);
+  }
+};
+
+TEST_F(QuantizePackTest, VerifyPack) {
+  auto status = QuantizeModel(&builder_, &model_, &error_reporter_);
+
+  ASSERT_EQ(kTfLiteOk, status);
+
+  const auto subgraph = model_.subgraphs[0].get();
+
+  // The model should only have 3 inputs and 1 output.
+  EXPECT_EQ(subgraph->inputs.size(), 3);
+  EXPECT_EQ(subgraph->outputs.size(), 1);
+
+  const auto& op1 = subgraph->operators[1].get();
+  const auto& op2 = subgraph->operators[2].get();
+  const auto& op3 = subgraph->operators[3].get();
+  const auto& op4 = subgraph->operators[4].get();
+
+  ASSERT_EQ(model_.operator_codes[op1->opcode_index].get()->builtin_code,
+            BuiltinOperator_QUANTIZE);
+  ASSERT_EQ(model_.operator_codes[op2->opcode_index].get()->builtin_code,
+            BuiltinOperator_QUANTIZE);
+  ASSERT_EQ(model_.operator_codes[op3->opcode_index].get()->builtin_code,
+            BuiltinOperator_PACK);
+  ASSERT_EQ(model_.operator_codes[op4->opcode_index].get()->builtin_code,
+            BuiltinOperator_DEQUANTIZE);
+
+  const auto& pack_input0 = subgraph->tensors[op3->inputs[0]].get();
+  const auto& pack_input1 = subgraph->tensors[op3->inputs[1]].get();
+  const auto& pack_input2 = subgraph->tensors[op3->inputs[2]].get();
+
+  const auto& pack_output = subgraph->tensors[op3->outputs[0]].get();
+
+  // Check quantization parameters for input and output.
+  EXPECT_FLOAT_EQ(pack_input0->quantization->scale[0],
+                  pack_input1->quantization->scale[0]);
+  EXPECT_FLOAT_EQ(pack_input1->quantization->scale[0],
+                  pack_input2->quantization->scale[0]);
+  EXPECT_FLOAT_EQ(pack_input0->quantization->zero_point[0],
+                  pack_input1->quantization->zero_point[0]);
+  EXPECT_FLOAT_EQ(pack_input1->quantization->zero_point[0],
+                  pack_input2->quantization->zero_point[0]);
+
+  EXPECT_FLOAT_EQ(pack_input1->quantization->scale[0],
+                  pack_output->quantization->scale[0]);
+  EXPECT_FLOAT_EQ(pack_input1->quantization->zero_point[0],
+                  pack_output->quantization->zero_point[0]);
+
+  // Check type of input and output.
+  EXPECT_EQ(pack_output->type, TensorType_INT8);
+  EXPECT_EQ(pack_input0->type, TensorType_INT8);
+  EXPECT_EQ(pack_input1->type, TensorType_INT8);
+  EXPECT_EQ(pack_input2->type, TensorType_INT8);
+}
+
 class QuantizeMinimumMaximumTest
     : public QuantizeModelTest,
       public testing::WithParamInterface<const char*> {

tensorflow/lite/tools/optimize/test_util.cc

@@ -51,6 +51,8 @@ const char* kModelMixed = "mixed.bin";
 
 const char* kModelSplit = "split.bin";
 
+const char* kModelPack = "pack.bin";
+
 const char* kLstmCalibrated = "lstm_calibrated.bin";
 const char* kLstmQuantized = "lstm_quantized.bin";
 

tensorflow/lite/tools/optimize/test_util.h

@@ -79,6 +79,9 @@ extern const char* kModelMixed;
 // Test model with split op.
 extern const char* kModelSplit;
 
+// Test model with pack op.
+extern const char* kModelPack;
+
 // Test model with LSTM op that has layer norm, has projection, without
 // peephole, without cifg.
 extern const char* kLstmCalibrated;