From 82e32167b625051363abc37ac6af8d17882bb11b Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Thu, 21 Nov 2019 09:30:25 -0800
Subject: [PATCH] Enable affine quantized tensor in writer_lib.

PiperOrigin-RevId: 281766372
Change-Id: I65d2dacdf4ea2a5df61406abd0b35bdffda2571e
---
 .../lite/experimental/writer/writer_lib.cc    | 37 +++++--------------
 1 file changed, 9 insertions(+), 28 deletions(-)

diff --git a/tensorflow/lite/experimental/writer/writer_lib.cc b/tensorflow/lite/experimental/writer/writer_lib.cc
index bb0d2fc5188..a4828fbc537 100644
--- a/tensorflow/lite/experimental/writer/writer_lib.cc
+++ b/tensorflow/lite/experimental/writer/writer_lib.cc
@@ -164,38 +164,19 @@ InterpreterWriter::ExportTensors(flatbuffers::FlatBufferBuilder* fbb) {
       // Primitive type.
       TensorType type = TfLiteTypeToSchemaType(tensor->type);
       // Handle quantization
-      flatbuffers::Offset<QuantizationParameters> quantization_params;
-
       const flatbuffers::Offset<flatbuffers::Vector<float>> null_array;
       flatbuffers::Offset<flatbuffers::Vector<float>> scale_array;
       flatbuffers::Offset<flatbuffers::Vector<int64_t>> zero_point_array;
-      // Multi channel quantization.
-      if (tensor->quantization.type == kTfLiteAffineQuantization) {
-        const TfLiteAffineQuantization* params =
-            reinterpret_cast<TfLiteAffineQuantization*>(
-                tensor->quantization.params);
-        const size_t num_scales = params->scale->size;
-
-        const int channel_index = params->quantized_dimension;
-        std::vector<float> scale_vector(
-            {params->scale->data, params->scale->data + num_scales});
-        std::vector<int64_t> zero_point_vector(
-            {params->zero_point->data, params->zero_point->data + num_scales});
-        scale_array = fbb->CreateVector<float>(scale_vector);
-        zero_point_array = fbb->CreateVector<int64_t>(zero_point_vector);
-        quantization_params = CreateQuantizationParameters(
-            *fbb, null_array, null_array, scale_array, zero_point_array,
-            QuantizationDetails_NONE, 0, channel_index);
-      } else {
-        // Quantization with a single argument array.
-        if (tensor->params.scale != 0.f) {
-          scale_array = fbb->CreateVector<float>({tensor->params.scale});
-          zero_point_array =
-              fbb->CreateVector<int64_t>({tensor->params.zero_point});
-        }
-        quantization_params = CreateQuantizationParameters(
-            *fbb, null_array, null_array, scale_array, zero_point_array);
+      if (tensor->params.scale != 0.f) {
+        // We have quantization, make a single arugment array (multi channel
+        // quant needs updating here).
+        scale_array = fbb->CreateVector<float>({tensor->params.scale});
+        zero_point_array =
+            fbb->CreateVector<int64_t>({tensor->params.zero_point});
       }
+      flatbuffers::Offset<QuantizationParameters> quantization_params =
+          CreateQuantizationParameters(*fbb, null_array, null_array,
+                                       scale_array, zero_point_array);
       // Shape
       TfLiteIntArrayView shape_view(tensor->dims);
       std::vector<int> shape =