Adding support for const ops that "broadcast" their own values.

PiperOrigin-RevId: 181345661

diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
index 0885c6eddbea41aa5d97f6b34f03b2ffe501e2de..a63060c3408a38a91f7e69475c34e729004a94e7 100644 (file)
--- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
+++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
@@ -55,10 +55,6 @@ std::map<string, string> kBrokenTests = {
     {R"(constant.*int32)", "68808744"},
     {R"(mul.*int32)", "68808744"},
 
-    // Toco or TFLite has a bug to deal with some constant functions with
-    // more than 1 element.
-    {R"(constant.*input_shape=\[(2|2,2,2,2)\])", "68721522"},
-
     // Pad only supports 4D float32 tensors.
     {R"(paddtype=.*,input_shape=\[.,.\],paddings=\[\[.,.\],\[.,.\]\])",
      "70527055"},

diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index 983b93aa3b6a67c466d6c2a73c949075f20651dd..0cec10e7d3c80a77228af3cfe12d5a3e0a7789f2 100644 (file)
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -173,7 +173,11 @@ void ImportFloatArray(const TensorProto& input_tensor, Array* output_array) {
   auto& output_float_data =
       output_array->GetMutableBuffer<ArrayDataType::kFloat>().data;
   output_float_data.resize(input_flat_size);
-  if (input_tensor.float_val_size()) {
+  if (input_tensor.float_val_size() == 1) {
+    for (int i = 0; i < input_flat_size; i++) {
+      output_float_data[i] = input_tensor.float_val(0);
+    }
+  } else if (input_tensor.float_val_size() == input_flat_size) {
     for (int i = 0; i < input_tensor.float_val_size(); i++) {
       output_float_data[i] = input_tensor.float_val(i);
     }