Enable bool input for CAST (#2401)

BUG=none
tensorflow · Jan 22, 2024 · 324ae1e · 324ae1e
1 parent 4b2bdd1
commit 324ae1e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/tensorflow/lite/micro/kernels/cast.cc b/tensorflow/lite/micro/kernels/cast.cc
@@ -101,6 +101,9 @@ TfLiteStatus CastEval(TfLiteContext* context, TfLiteNode* node) {
     case kTfLiteFloat32:
       return copyToTensor(context, tflite::micro::GetTensorData<float>(input),
                           output, num_elements);
+    case kTfLiteBool:
+      return copyToTensor(context, tflite::micro::GetTensorData<bool>(input),
+                          output, num_elements);
     default:
       // Unsupported type.
       MicroPrintf("Input type %s (%d) not supported.",

diff --git a/tensorflow/lite/micro/kernels/cast_test.cc b/tensorflow/lite/micro/kernels/cast_test.cc
@@ -129,4 +129,12 @@ TF_LITE_MICRO_TEST(CastUInt32ToInt32) {
   tflite::testing::TestCast(input_dims, input_values, golden, output_data);
 }
 
+TF_LITE_MICRO_TEST(CastBoolToFloat) {
+  float output_data[6];
+  int input_dims[] = {2, 2, 3};
+  const bool input_values[] = {true, true, false, true, false, true};
+  const float golden[] = {1.f, 1.0f, 0.f, 1.0f, 0.0f, 1.0f};
+  tflite::testing::TestCast(input_dims, input_values, golden, output_data);
+}
+
 TF_LITE_MICRO_TESTS_END