Fix flexbuffers vector creation bug in MLIR TFLite export

---

When exporting StableHLOComposite attributes, flexbuffers::Builder was being called with the attribute name as a key for every vector element. In flexbuffers, passing a key to a Vector element pushes both the key and the value to the stack, causing the vector to be parsed with twice its intended length (with the key interleaved).

This change removes the key from the vector element calls, and only uses the key when starting the parent vector, which matches the flexbuffers::Builder API expectation for Vectors.

- Also fixed a typo where `isa<StringAttr>(attr)` was checking the Array attribute instead of the element.

- Updated test to pass actual vector values for serialization

PiperOrigin-RevId: 941177131

Author: tensorflower-gardener

tensorflow/compiler/mlir/lite/flatbuffer_export.cc

@@ -1904,39 +1904,35 @@ uint32_t Translator::GetOpcodeIndex(const std::string& op_name,
 
 void CreateFlexbufferVector(
     const std::unique_ptr<flexbuffers::Builder>& flex_builder,
-    std::string& name, const mlir::Attribute& attr) {
-  auto start = flex_builder->StartVector(name.c_str());
+    std::optional<absl::string_view> key, const mlir::Attribute& attr) {
+  auto start = key.has_value()
+                   ? flex_builder->StartVector(std::string(*key).c_str())
+                   : flex_builder->StartVector();
   auto array = mlir::cast<mlir::vhlo::ArrayV1Attr>(attr).getValue();
 
   for (int i = 0; i < array.size(); i++) {
     if (llvm::isa<mlir::BoolAttr>(array[i])) {
-      flex_builder->Bool(name.c_str(),
-                         mlir::cast<mlir::BoolAttr>(array[i]).getValue());
-    } else if (llvm::isa<mlir::StringAttr>(attr)) {
+      flex_builder->Bool(mlir::cast<mlir::BoolAttr>(array[i]).getValue());
+    } else if (llvm::isa<mlir::StringAttr>(array[i])) {
       flex_builder->String(
-          name.c_str(),
           mlir::cast<mlir::StringAttr>(array[i]).getValue().str());
     } else if (llvm::isa<mlir::vhlo::BooleanV1Attr>(array[i])) {
       flex_builder->Bool(
-          name.c_str(),
           mlir::cast<mlir::vhlo::BooleanV1Attr>(array[i]).getValue());
     } else if (llvm::isa<mlir::vhlo::StringV1Attr>(array[i])) {
       flex_builder->String(
-          name.c_str(),
           mlir::cast<mlir::vhlo::StringV1Attr>(array[i]).getValue().str());
     } else if (llvm::isa<mlir::vhlo::IntegerV1Attr>(array[i])) {
-      flex_builder->Int(name.c_str(),
-                        mlir::cast<mlir::vhlo::IntegerV1Attr>(array[i])
+      flex_builder->Int(mlir::cast<mlir::vhlo::IntegerV1Attr>(array[i])
                             .getValue()
                             .getSExtValue());
     } else if (llvm::isa<mlir::vhlo::FloatV1Attr>(array[i])) {
-      flex_builder->Float(name.c_str(),
-                          mlir::cast<mlir::vhlo::FloatV1Attr>(array[i])
+      flex_builder->Float(mlir::cast<mlir::vhlo::FloatV1Attr>(array[i])
                               .getValue()
                               .convertToFloat());
 
     } else if (llvm::isa<mlir::vhlo::ArrayV1Attr>(array[i])) {
-      CreateFlexbufferVector(flex_builder, name, array[i]);
+      CreateFlexbufferVector(flex_builder, std::nullopt, array[i]);
     }
   }
 

tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/composite.mlir

@@ -88,7 +88,7 @@
 // CHECK-NEXT:      builtin_options_2: {
 // CHECK-NEXT:        name: "test.TEST_COMPOSITE",
 // CHECK-NEXT:        decomposition_subgraph_index: 2,
-// CHECK-NEXT:        composite_attributes: [ 0, 0, 1, 0, 0, 36, 1 ]
+// CHECK-NEXT:        composite_attributes: [ 109, 121, 95, 97, 114, 114, 97, 121, 0, 1, 97, 0, 1, 98, 0, 2, 6, 4, 20, 20, 1, 21, 1, 1, 1, 9, 40, 2, 36, 1 ]
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }, {
 // CHECK-NEXT:      inputs: [ 4, 1 ],
@@ -97,7 +97,7 @@
 // CHECK-NEXT:      builtin_options_2: {
 // CHECK-NEXT:        name: "test.TEST_COMPOSITE",
 // CHECK-NEXT:        decomposition_subgraph_index: 1,
-// CHECK-NEXT:        composite_attributes: [ 0, 0, 1, 0, 0, 36, 1 ]
+// CHECK-NEXT:        composite_attributes: [ 109, 121, 95, 97, 114, 114, 97, 121, 0, 1, 97, 0, 1, 98, 0, 2, 6, 4, 20, 20, 1, 21, 1, 1, 1, 9, 40, 2, 36, 1 ]
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }, {
 // CHECK-NEXT:      opcode_index: 1,
@@ -294,8 +294,8 @@
 func.func @main(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> (tensor<10xf32>) {
   %cst = arith.constant dense<1.000000e+01> : tensor<f32>
   %cst_0 = arith.constant dense<2.000000e+01> : tensor<f32>
-  %0 = "vhlo.composite_v1"(%arg0, %arg1) <{composite_attributes = #vhlo.dict_v1<{}>, decomposition = #vhlo.string_v1<"XlaCallModule_test.TEST_COMPOSITE.impl_0_0">, name = #vhlo.string_v1<"test.TEST_COMPOSITE">, version = #vhlo.integer_v1<0 : i64>}> : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
-  %1 = "vhlo.composite_v1"(%0, %arg1) <{composite_attributes = #vhlo.dict_v1<{}>, decomposition = #vhlo.string_v1<"XlaCallModule_test.TEST_COMPOSITE.impl_0">, name = #vhlo.string_v1<"test.TEST_COMPOSITE">, version = #vhlo.integer_v1<0 : i64>}> : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
+  %0 = "vhlo.composite_v1"(%arg0, %arg1) <{composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_array"> = #vhlo.array_v1<[#vhlo.string_v1<"a">, #vhlo.string_v1<"b">]>}>, decomposition = #vhlo.string_v1<"XlaCallModule_test.TEST_COMPOSITE.impl_0_0">, name = #vhlo.string_v1<"test.TEST_COMPOSITE">, version = #vhlo.integer_v1<0 : i64>}> : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
+  %1 = "vhlo.composite_v1"(%0, %arg1) <{composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_array"> = #vhlo.array_v1<[#vhlo.string_v1<"a">, #vhlo.string_v1<"b">]>}>, decomposition = #vhlo.string_v1<"XlaCallModule_test.TEST_COMPOSITE.impl_0">, name = #vhlo.string_v1<"test.TEST_COMPOSITE">, version = #vhlo.integer_v1<0 : i64>}> : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
   %2 = tfl.add(%1, %cst) <{fused_activation_function = "NONE"}> : (tensor<10xf32>, tensor<f32>) -> tensor<10xf32>
   %3 = tfl.sub(%2, %cst_0) <{fused_activation_function = "NONE"}> : (tensor<10xf32>, tensor<f32>) -> tensor<10xf32>
   return %3 : tensor<10xf32>