Modifying tensor_to_vector_float_calculator to take in D_BFLOAT16 values

PiperOrigin-RevId: 565189254

mediapipe/calculators/tensorflow/BUILD

@@ -723,6 +723,7 @@ cc_library(
         "//mediapipe/framework:calculator_framework",
         "//mediapipe/framework/port:ret_check",
         "//mediapipe/framework/port:status",
+        "@org_tensorflow//tensorflow/core/platform:bfloat16",
     ] + select({
         "//conditions:default": [
             "@org_tensorflow//tensorflow/core:framework",
@@ -1139,6 +1140,7 @@ cc_test(
         "//mediapipe/util:packet_test_util",
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
+        "@org_tensorflow//tensorflow/core/platform:bfloat16",
     ],
 )
 

mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator.cc

@@ -15,12 +15,16 @@
 // Calculator converts from one-dimensional Tensor of DT_FLOAT to vector<float>
 // OR from (batched) two-dimensional Tensor of DT_FLOAT to vector<vector<float>.
 
+#include <memory>
+#include <vector>
+
 #include "mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_options.pb.h"
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/port/ret_check.h"
 #include "mediapipe/framework/port/status.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/bfloat16.h"
 
 namespace mediapipe {
 
@@ -76,23 +80,33 @@ absl::Status TensorToVectorFloatCalculator::Open(CalculatorContext* cc) {
 absl::Status TensorToVectorFloatCalculator::Process(CalculatorContext* cc) {
   const tf::Tensor& input_tensor =
       cc->Inputs().Index(0).Value().Get<tf::Tensor>();
-  RET_CHECK(tf::DT_FLOAT == input_tensor.dtype())
+  RET_CHECK(tf::DT_FLOAT == input_tensor.dtype() ||
-      << "expected DT_FLOAT input but got "
+            tf::DT_BFLOAT16 == input_tensor.dtype())
+      << "expected DT_FLOAT or DT_BFLOAT_16 input but got "
       << tensorflow::DataTypeString(input_tensor.dtype());
 
   if (options_.tensor_is_2d()) {
     RET_CHECK(2 == input_tensor.dims())
         << "Expected 2-dimensional Tensor, but the tensor shape is: "
         << input_tensor.shape().DebugString();
-    auto output = absl::make_unique<std::vector<std::vector<float>>>(
+    auto output = std::make_unique<std::vector<std::vector<float>>>(
         input_tensor.dim_size(0), std::vector<float>(input_tensor.dim_size(1)));
     for (int i = 0; i < input_tensor.dim_size(0); ++i) {
       auto& instance_output = output->at(i);
-      const auto& slice = input_tensor.Slice(i, i + 1).unaligned_flat<float>();
+      if (tf::DT_BFLOAT16 == input_tensor.dtype()) {
+        const auto& slice =
+            input_tensor.Slice(i, i + 1).unaligned_flat<tf::bfloat16>();
+        for (int j = 0; j < input_tensor.dim_size(1); ++j) {
+          instance_output.at(j) = static_cast<float>(slice(j));
+        }
+      } else {
+        const auto& slice =
+            input_tensor.Slice(i, i + 1).unaligned_flat<float>();
         for (int j = 0; j < input_tensor.dim_size(1); ++j) {
           instance_output.at(j) = slice(j);
         }
       }
+    }
     cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
   } else {
     if (!options_.flatten_nd()) {
@@ -101,11 +115,18 @@ absl::Status TensorToVectorFloatCalculator::Process(CalculatorContext* cc) {
           << "tensor shape is: " << input_tensor.shape().DebugString();
     }
     auto output =
-        absl::make_unique<std::vector<float>>(input_tensor.NumElements());
+        std::make_unique<std::vector<float>>(input_tensor.NumElements());
+    if (tf::DT_BFLOAT16 == input_tensor.dtype()) {
+      const auto& tensor_values = input_tensor.unaligned_flat<tf::bfloat16>();
+      for (int i = 0; i < input_tensor.NumElements(); ++i) {
+        output->at(i) = static_cast<float>(tensor_values(i));
+      }
+    } else {
       const auto& tensor_values = input_tensor.unaligned_flat<float>();
       for (int i = 0; i < input_tensor.NumElements(); ++i) {
         output->at(i) = tensor_values(i);
       }
+    }
     cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
   }
 

mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_test.cc

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <memory>
+
 #include "mediapipe/calculators/tensorflow/tensor_to_vector_float_calculator_options.pb.h"
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/calculator_runner.h"
@@ -19,6 +21,7 @@
 #include "mediapipe/util/packet_test_util.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/bfloat16.h"
 
 namespace mediapipe {
 
@@ -72,6 +75,62 @@ TEST_F(TensorToVectorFloatCalculatorTest, ConvertsToVectorFloat) {
   }
 }
 
+TEST_F(TensorToVectorFloatCalculatorTest, CheckBFloat16Type) {
+  SetUpRunner(false, false);
+  const tf::TensorShape tensor_shape(std::vector<tf::int64>{5});
+  auto tensor = std::make_unique<tf::Tensor>(tf::DT_BFLOAT16, tensor_shape);
+  auto tensor_vec = tensor->vec<tf::bfloat16>();
+  for (int i = 0; i < 5; ++i) {
+    tensor_vec(i) = static_cast<tf::bfloat16>(1 << i);
+  }
+
+  const int64_t time = 1234;
+  runner_->MutableInputs()->Index(0).packets.push_back(
+      Adopt(tensor.release()).At(Timestamp(time)));
+
+  EXPECT_TRUE(runner_->Run().ok());
+  const std::vector<Packet>& output_packets =
+      runner_->Outputs().Index(0).packets;
+  EXPECT_EQ(1, output_packets.size());
+  EXPECT_EQ(time, output_packets[0].Timestamp().Value());
+  const std::vector<float>& output_vector =
+      output_packets[0].Get<std::vector<float>>();
+
+  EXPECT_EQ(5, output_vector.size());
+  for (int i = 0; i < 5; ++i) {
+    const float expected = static_cast<float>(1 << i);
+    EXPECT_EQ(expected, output_vector[i]);
+  }
+}
+
+TEST_F(TensorToVectorFloatCalculatorTest, CheckBFloat16TypeAllDim) {
+  SetUpRunner(false, true);
+  const tf::TensorShape tensor_shape(std::vector<tf::int64>{2, 2, 2});
+  auto tensor = std::make_unique<tf::Tensor>(tf::DT_BFLOAT16, tensor_shape);
+  auto slice = tensor->flat<tf::bfloat16>();
+  for (int i = 0; i < 2 * 2 * 2; ++i) {
+    // 2^i can be represented exactly in floating point numbers if 'i' is small.
+    slice(i) = static_cast<tf::bfloat16>(1 << i);
+  }
+
+  const int64_t time = 1234;
+  runner_->MutableInputs()->Index(0).packets.push_back(
+      Adopt(tensor.release()).At(Timestamp(time)));
+
+  EXPECT_TRUE(runner_->Run().ok());
+  const std::vector<Packet>& output_packets =
+      runner_->Outputs().Index(0).packets;
+  EXPECT_EQ(1, output_packets.size());
+  EXPECT_EQ(time, output_packets[0].Timestamp().Value());
+  const std::vector<float>& output_vector =
+      output_packets[0].Get<std::vector<float>>();
+  EXPECT_EQ(2 * 2 * 2, output_vector.size());
+  for (int i = 0; i < 2 * 2 * 2; ++i) {
+    const float expected = static_cast<float>(1 << i);
+    EXPECT_EQ(expected, output_vector[i]);
+  }
+}
+
 TEST_F(TensorToVectorFloatCalculatorTest, ConvertsBatchedToVectorVectorFloat) {
   SetUpRunner(true, false);
   const tf::TensorShape tensor_shape(std::vector<tf::int64>{1, 5});