commit 27c820fed31089b3c1857022b6d143575adb782a
author Alex Van Damme <atv@google.com> Wed May 29 13:28:46 2024 -0700
committer Alex Van Damme <atv@google.com> Wed May 29 13:28:46 2024 -0700
tree eef76e200da96d5bfd657f2f3c29b44d768d69f7
parent aa3fb6a13c9d8bbc8de57420c6c27ccd1f52e9f2

Add int16 pack and transpose

Change-Id: I8d63eadd27cb20868045666efb3b0369efc2c1cb

tensorflow/lite/micro/kernels/pack.cc

diff --git a/tensorflow/lite/micro/kernels/pack.cc b/tensorflow/lite/micro/kernels/pack.cc
index 7b4aeef..5ee2759 100644
--- a/tensorflow/lite/micro/kernels/pack.cc
+++ b/tensorflow/lite/micro/kernels/pack.cc
@@ -85,6 +85,9 @@
       return PackImpl<int8_t>(context, node, output, data->values_count,
                               data->axis);
     }
+    case kTfLiteInt16: {
+      return PackImpl<int16_t>(context, node, output, data->values_count, data->axis);
+    }
     case kTfLiteInt32: {
       return PackImpl<int32_t>(context, node, output, data->values_count,
                                data->axis);

tensorflow/lite/micro/kernels/transpose.cc

diff --git a/tensorflow/lite/micro/kernels/transpose.cc b/tensorflow/lite/micro/kernels/transpose.cc
index 710bfca..c57812b 100644
--- a/tensorflow/lite/micro/kernels/transpose.cc
+++ b/tensorflow/lite/micro/kernels/transpose.cc
@@ -97,6 +97,12 @@
                                tflite::micro::GetTensorShape(output),
                                tflite::micro::GetTensorData<float>(output));
       break;
+    case kTfLiteInt16:
+      reference_ops::Transpose(params, tflite::micro::GetTensorShape(input),
+          tflite::micro::GetTensorData<int16_t>(input),
+          tflite::micro::GetTensorShape(output),
+          tflite::micro::GetTensorData<int16_t>(output));
+      break;
     case kTfLiteInt8:
       reference_ops::Transpose(params, tflite::micro::GetTensorShape(input),
                                tflite::micro::GetTensorData<int8_t>(input),