Sync from upstream TF.

Author: TFLM-bot

tensorflow/compiler/mlir/lite/schema/schema_utils.cc

@@ -15,8 +15,12 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
 
 #include <algorithm>
+#include <complex>
+#include <cstddef>
+#include <cstdint>
 
 #include "tensorflow/compiler/mlir/lite/kernels/internal/compatibility_macros.h"
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
 
 namespace tflite {
 
@@ -59,4 +63,51 @@ BuiltinOperator GetBuiltinCode(const OperatorCodeT* op_code) {
                                              op_code->deprecated_builtin_code));
 }
 
+size_t TensorTypeGetSize(::tflite::TensorType data_type) {
+  switch (data_type) {
+    case ::tflite::TensorType_FLOAT32:
+      static_assert(sizeof(float) == 4, "");
+      return 4;
+    case ::tflite::TensorType_FLOAT16:
+      static_assert(sizeof(int16_t) == 2, "");
+      return 2;
+    case ::tflite::TensorType_INT32:
+      static_assert(sizeof(int32_t) == 4, "");
+      return 4;
+    case ::tflite::TensorType_UINT8:
+      static_assert(sizeof(uint8_t) == 1, "");
+      return 1;
+    case ::tflite::TensorType_INT64:
+      static_assert(sizeof(int64_t) == 8, "");
+      return 8;
+    case ::tflite::TensorType_BOOL:
+      return sizeof(bool);
+    case ::tflite::TensorType_INT16:
+      static_assert(sizeof(int16_t) == 2, "");
+      return 2;
+    case ::tflite::TensorType_COMPLEX64:
+      static_assert(sizeof(std::complex<float>) == 8, "");
+      return 8;
+    case ::tflite::TensorType_INT8:
+      static_assert(sizeof(int8_t) == 1, "");
+      return 1;
+    case ::tflite::TensorType_FLOAT64:
+      static_assert(sizeof(double) == 8, "");
+      return 8;
+    case ::tflite::TensorType_COMPLEX128:
+      static_assert(sizeof(std::complex<double>) == 16, "");
+      return 16;
+    case ::tflite::TensorType_UINT64:
+      static_assert(sizeof(uint64_t) == 8, "");
+      return 8;
+    case ::tflite::TensorType_UINT32:
+      static_assert(sizeof(uint32_t) == 4, "");
+      return 4;
+    case ::tflite::TensorType_UINT16:
+      static_assert(sizeof(uint16_t) == 2, "");
+      return 2;
+    default:
+      return 0;
+  }
+}
 }  // namespace tflite

tensorflow/compiler/mlir/lite/schema/schema_utils.h

@@ -15,6 +15,8 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_MLIR_LITE_SCHEMA_SCHEMA_UTILS_H_
 #define TENSORFLOW_COMPILER_MLIR_LITE_SCHEMA_SCHEMA_UTILS_H_
 
+#include <cstddef>
+
 #include "flatbuffers/flatbuffers.h"
 #include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
 
@@ -28,6 +30,11 @@ BuiltinOperator GetBuiltinCode(const OperatorCode *op_code);
 
 BuiltinOperator GetBuiltinCode(const OperatorCodeT *op_code);
 
+// Returns the size of the given TensorType in bytes, or 0 if the TensorType is
+// not supported, this function should be aligned with TfLiteTypeGetSize in
+// lite/kernels/kernel_util.h.
+size_t TensorTypeGetSize(::tflite::TensorType data_type);
+
 }  // namespace tflite
 
 #endif  // TENSORFLOW_COMPILER_MLIR_LITE_SCHEMA_SCHEMA_UTILS_H_

tensorflow/lite/kernels/internal/reference/broadcast_to.h

@@ -15,6 +15,8 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_BROADCAST_TO_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_BROADCAST_TO_H_
 
+#include <cstddef>
+
 #include "tensorflow/lite/kernels/internal/common.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
@@ -83,7 +85,8 @@ inline void BroadcastTo(const RuntimeShape& unextended_input_shape,
   // If non-broadcasting, just copy data from input to output tensor.
   if (last_broadcast_dim == -1) {
     memcpy(output_data, input_data,
-           unextended_input_shape.FlatSize() * TfLiteTypeGetSize(data_type));
+           static_cast<size_t>(unextended_input_shape.FlatSize()) *
+               static_cast<size_t>(TfLiteTypeGetSize(data_type)));
     return;
   }