[Unique-10] ConvertPrecision transformation and new attribute for output type (#14229)

* Add attribute for last output element type

* Add convert precision transformation and tests

* Update Unique python API with new attribute

* Update Unique-10 op specification

* Update docstrings

* Update visitor tests

* Add type prop tests for the new attribute

* Check axis constant before shape, compare with partial shape

docs/ops/movement/Unique_10.md

@@ -29,6 +29,14 @@ The operator can either work in elementwise mode searching for unique values in
   * **Default value**: "i64"
   * **Required**: *no*
 
+* *count_element_type*
+
+  * **Description**: controls the data type of the last output tensor.
+  * **Range of values**: "i64" or "i32"
+  * **Type**: string
+  * **Default value**: "i64"
+  * **Required**: *no*
+
 **Inputs**
 
 * **1**: A tensor of type *T* and arbitrary shape. **Required.**
@@ -40,7 +48,7 @@ The operator can either work in elementwise mode searching for unique values in
 * **1**: The output tensor containing unique elements (individual values or subtensors). This tensor's type matches the type of the first input tensor: *T*. The values in this tensor are either sorted ascendingly or maintain the same order as in the input tensor. The shape of this output depends on the values of the input tensor and will very often be dynamic. Please refer to the article describing how [Dynamic Shapes](https://docs.openvino.ai/latest/openvino_docs_OV_UG_DynamicShapes.html) are handled in OpenVINO.
 * **2**: The output tensor containing indices of the locations of unique elements. The indices map the elements in the first output tensor to their locations in the input tensor. The index always points to the first occurrence of a given unique output element in the input tensor. This is a 1D tensor with type controlled by the `index_element_type` attribute.
 * **3**: The output tensor containing indices of the locations of elements of the input tensor in the first output tensor. This means that for each element of the input tensor this output will point to the unique value in the first output tensor of this operator. This is a 1D tensor with type controlled by the `index_element_type` attribute.
-* **4**: The output tensor containing the number of occurrences of each unique value produced by this operator in the first output tensor. This is a 1D tensor with type `int64`.
+* **4**: The output tensor containing the number of occurrences of each unique value produced by this operator in the first output tensor. This is a 1D tensor with type controlled by the `count_element_type` attribute.
 
 **Types**
 
@@ -74,7 +82,7 @@ The operator can either work in elementwise mode searching for unique values in
             <dim>-1</dim>
         </port>
         <port id="4" precision="I32">
-            <dim>-1</dim>
+            <dim>3</dim>
         </port>
         <port id="5" precision="I64">
             <dim>-1</dim>
@@ -96,11 +104,11 @@ The operator can either work in elementwise mode searching for unique values in
         <port id="1" precision="FP32">
             <dim>-1</dim>
         </port>
-        <port id="2" precision="I32">
+        <port id="2" precision="I64">
             <dim>-1</dim>
         </port>
-        <port id="3" precision="I32">
-            <dim>-1</dim>
+        <port id="3" precision="I64">
+            <dim>9</dim>
         </port>
         <port id="4" precision="I64">
             <dim>-1</dim>
@@ -108,3 +116,30 @@ The operator can either work in elementwise mode searching for unique values in
     </output>
 </layer>
 ```
+
+*Example 3: no axis provided, non-default outputs precision *
+```xml
+<layer ... type="Unique" ... >
+    <data sorted="false" index_element_type="i32" count_element_type="i32"/>
+    <input>
+        <port id="0" precision="FP32">
+            <dim>3</dim>
+            <dim>3</dim>
+        </port>
+    </input>
+    <output>
+        <port id="1" precision="FP32">
+            <dim>-1</dim>
+        </port>
+        <port id="2" precision="I32">
+            <dim>-1</dim>
+        </port>
+        <port id="3" precision="I32">
+            <dim>9</dim>
+        </port>
+        <port id="4" precision="I32">
+            <dim>-1</dim>
+        </port>
+    </output>
+</layer>
+```

src/bindings/python/src/compatibility/ngraph/opset10/ops.py

@@ -141,6 +141,7 @@ def unique(
     axis: Optional[NodeInput] = None,
     sorted: Optional[bool] = True,
     index_element_type: Optional[str] = "i64",
+    count_element_type: Optional[str] = "i64",
     name: Optional[str] = None,
 ) -> Node:
     """Operator which selects and returns unique elements or unique slices of the input tensor.
@@ -154,6 +155,8 @@ def unique(
                                 Default value: True.
     :param  index_element_type: (Optional) The data type set for outputs containing indices.
                                 Default value: "i64".
+    :param  count_element_type: (Optional) The data type set for the output with repetition count.
+                                Default value: "i64".
     :param name:                (Optional) A name for the output node. Default value: None.
     :return: Node representing Unique operation.
     """
@@ -165,5 +168,6 @@ def unique(
     attributes = {
         "sorted": sorted,
         "index_element_type": index_element_type,
+        "count_element_type": count_element_type,
     }
     return _get_node_factory_opset10().create("Unique", inputs, attributes)

src/bindings/python/src/openvino/runtime/opset10/ops.py

@@ -141,6 +141,7 @@ def unique(
     axis: Optional[NodeInput] = None,
     sorted: Optional[bool] = True,
     index_element_type: Optional[str] = "i64",
+    count_element_type: Optional[str] = "i64",
     name: Optional[str] = None,
 ) -> Node:
     """Operator which selects and returns unique elements or unique slices of the input tensor.
@@ -154,6 +155,8 @@ def unique(
                                 Default value: True.
     :param  index_element_type: (Optional) The data type set for outputs containing indices.
                                 Default value: "i64".
+    :param  count_element_type: (Optional) The data type set for the output with repetition count.
+                                Default value: "i64".
     :param name:                (Optional) A name for the output node. Default value: None.
     :return: Node representing Unique operation.
     """
@@ -165,5 +168,6 @@ def unique(
     attributes = {
         "sorted": sorted,
         "index_element_type": index_element_type,
+        "count_element_type": count_element_type,
     }
     return _get_node_factory_opset10().create("Unique", inputs, attributes)

src/bindings/python/tests/test_graph/test_create_op.py

@@ -2294,7 +2294,7 @@ def test_unique_opset10():
     assert node.get_output_element_type(3) == Type.i64
 
     # Axis default, means flattened result
-    node = ov_opset10.unique(input_node, None, False, "i32")
+    node = ov_opset10.unique(input_node, None, False, "i32", "i32")
 
     assert node.get_type_name() == "Unique"
     assert node.get_sorted() is False
@@ -2308,7 +2308,7 @@ def test_unique_opset10():
     assert node.get_output_element_type(0) == Type.f32
     assert node.get_output_element_type(1) == Type.i32
     assert node.get_output_element_type(2) == Type.i32
-    assert node.get_output_element_type(3) == Type.i64
+    assert node.get_output_element_type(3) == Type.i32
 
     # All arguments default
     node = ov_opset10.unique(input_node)

src/bindings/python/tests_compatibility/test_ngraph/test_create_op.py

@@ -2359,7 +2359,7 @@ def test_unique_opset10():
     assert node.get_output_element_type(3) == Type.i64
 
     # Axis default, means flattened result
-    node = ng_opset10.unique(input_node, None, False, "i32")
+    node = ng_opset10.unique(input_node, None, False, "i32", "i32")
 
     assert node.get_type_name() == "Unique"
     assert node.get_sorted() is False
@@ -2373,7 +2373,7 @@ def test_unique_opset10():
     assert node.get_output_element_type(0) == Type.f32
     assert node.get_output_element_type(1) == Type.i32
     assert node.get_output_element_type(2) == Type.i32
-    assert node.get_output_element_type(3) == Type.i64
+    assert node.get_output_element_type(3) == Type.i32
 
     # All arguments default
     node = ng_opset10.unique(input_node)

src/common/transformations/src/transformations/convert_precision.cpp

@@ -7,6 +7,7 @@
 #include <memory>
 #include <ngraph/runtime/reference/convert.hpp>
 #include <openvino/opsets/opset1.hpp>
+#include <openvino/opsets/opset10.hpp>
 #include <openvino/opsets/opset3.hpp>
 #include <openvino/opsets/opset4.hpp>
 #include <openvino/opsets/opset5.hpp>
@@ -26,6 +27,7 @@ bool fuse_type_to_constant(const std::shared_ptr<ngraph::Node>& node,
 bool fuse_type_to_shapeof(const std::shared_ptr<ngraph::Node>& node, ngraph::element::Type to, size_t idx);
 bool fuse_type_to_shapeof_v0(const std::shared_ptr<ngraph::Node>& node, ngraph::element::Type to, size_t idx);
 bool fuse_type_to_random_uniform_v8(const std::shared_ptr<ngraph::Node>& node, ngraph::element::Type to, size_t idx);
+bool fuse_type_to_unique_v10(const std::shared_ptr<ngraph::Node>& node, ngraph::element::Type to, size_t idx);
 bool fuse_type_to_range_v4(const std::shared_ptr<ngraph::Node>& node, ngraph::element::Type to, size_t idx);
 bool fuse_type_to_parameter(const std::shared_ptr<ngraph::Node>& node, ngraph::element::Type to, size_t idx);
 bool fuse_type_to_convert(const std::shared_ptr<ngraph::Node>& node, ngraph::element::Type to, size_t idx);
@@ -288,6 +290,7 @@ bool ov::pass::ConvertPrecision::run_on_model(const std::shared_ptr<ngraph::Func
         {opset4::ReduceLogicalOr::get_type_info_static(), fuse_type_to_reduce_logical<opset4::ReduceLogicalOr>},
         {opset1::ShapeOf::get_type_info_static(), fuse_type_to_shapeof_v0},
         {opset4::Range::get_type_info_static(), fuse_type_to_range_v4},
+        {opset10::Unique::get_type_info_static(), fuse_type_to_unique_v10},
         {opset8::RandomUniform::get_type_info_static(), fuse_type_to_random_uniform_v8}};
 
     type_to_fuse.insert(m_additional_type_to_fuse_map.begin(), m_additional_type_to_fuse_map.end());
@@ -335,6 +338,20 @@ bool fuse_type_to_random_uniform_v8(const std::shared_ptr<ngraph::Node>& node, o
     return false;
 }
 
+bool fuse_type_to_unique_v10(const std::shared_ptr<Node>& node, ov::element::Type to, size_t idx) {
+    if (auto unique = ov::as_type_ptr<opset10::Unique>(node)) {
+        if (to == ov::element::i32 || to == ov::element::i64) {
+            if (idx == 1 || idx == 2) {
+                unique->set_index_element_type(to);
+            } else if (idx == 3) {
+                unique->set_count_element_type(to);
+            }
+        }
+    }
+    // No node replacement, so always return false
+    return false;
+}
+
 bool fuse_type_to_range_v4(const std::shared_ptr<ngraph::Node>& node, ov::element::Type to, size_t idx) {
     if (auto range = ov::as_type_ptr<opset4::Range>(node)) {
         if (to.is_integral_number() || to.is_real()) {

src/common/transformations/tests/utils/convert_precision.cpp

@@ -15,6 +15,7 @@
 #include <ngraph/opsets/opset4.hpp>
 #include <ngraph/opsets/opset5.hpp>
 #include <ngraph/opsets/opset8.hpp>
+#include <ngraph/opsets/opset10.hpp>
 #include <transformations/convert_precision.hpp>
 #include <transformations/utils/utils.hpp>
 #include <ngraph/pass/manager.hpp>
@@ -323,6 +324,37 @@ TEST(TransformationTests, ConvertPrecision_TopK) {
     ASSERT_FALSE(has_type<ngraph::element::Type_t::i64>(f));
 }
 
+TEST(TransformationTests, ConvertPrecision_Unique10) {
+    std::shared_ptr<ov::Model> model(nullptr);
+    {
+        auto input = std::make_shared<op::Parameter>(element::f16, Shape{15, 20, 3});
+        auto unique = std::make_shared<opset10::Unique>(input);
+
+        model = std::make_shared<ov::Model>(unique->outputs(), ParameterVector{input});
+
+        pass::Manager manager;
+
+        static const precisions_array precisions = {
+            { element::i64, element::i32 },
+            { element::f16, element::f32 }
+        };
+
+        manager.register_pass<pass::ConvertPrecision>(precisions);
+        manager.run_passes(model);
+    }
+
+    ASSERT_EQ(model->outputs().size(), 4);
+    EXPECT_EQ(model->outputs()[0].get_element_type(), element::f32);
+    EXPECT_EQ(model->outputs()[1].get_element_type(), element::i32);
+    EXPECT_EQ(model->outputs()[2].get_element_type(), element::i32);
+    EXPECT_EQ(model->outputs()[3].get_element_type(), element::i32);
+
+    EXPECT_EQ(model->get_results().size(), 4);
+
+    EXPECT_FALSE(has_type<element::Type_t::f16>(model));
+    EXPECT_FALSE(has_type<element::Type_t::i64>(model));
+}
+
 TEST(TransformationTests, ConvertPrecision_NonZero) {
     std::shared_ptr<Function> f(nullptr);
     {

src/core/include/openvino/op/unique.hpp

@@ -22,19 +22,25 @@ public:
     ///
     /// \param data     Input data tensor
     /// \param sorted   Controls the order of the returned unique values (sorts ascendingly when true)
-    /// \param index_element_type    The data type set for outputs containing indices
-    Unique(const Output<Node>& data, const bool sorted = true, const element::Type& index_element_type = element::i64);
+    /// \param index_element_type    The data type for outputs containing indices
+    /// \param count_element_type    The data type for output containing repetition count
+    Unique(const Output<Node>& data,
+           const bool sorted = true,
+           const element::Type& index_element_type = element::i64,
+           const element::Type& count_element_type = element::i64);
 
     /// \brief Constructs a Unique operation
     ///
     /// \param data     Input data tensor
     /// \param axis     An input tensor containing the axis value
     /// \param sorted   Controls the order of the returned unique values (sorts ascendingly when true)
-    /// \param index_element_type    The data type set for outputs containing indices
+    /// \param index_element_type    The data type for outputs containing indices
+    /// \param count_element_type    The data type for output containing repetition count
     Unique(const Output<Node>& data,
            const Output<Node>& axis,
            const bool sorted = true,
-           const element::Type& index_element_type = element::i64);
+           const element::Type& index_element_type = element::i64,
+           const element::Type& ount_element_type = element::i64);
     bool visit_attributes(AttributeVisitor& visitor) override;
 
     void validate_and_infer_types() override;
@@ -57,9 +63,18 @@ public:
         m_index_element_type = index_element_type;
     }
 
+    element::Type get_count_element_type() const {
+        return m_count_element_type;
+    }
+
+    void set_count_element_type(const element::Type& count_element_type) {
+        m_count_element_type = count_element_type;
+    }
+
 private:
     bool m_sorted = true;
     element::Type m_index_element_type = element::i64;
+    element::Type m_count_element_type = element::i64;
 };
 }  // namespace v10
 }  // namespace op

src/core/src/op/unique.cpp

@@ -93,20 +93,26 @@ std::tuple<Shape, Shape, Shape> calculate_static_output_shapes(const Tensor& inp
 }
 }  // namespace
 
-op::v10::Unique::Unique(const Output<Node>& data, const bool sorted, const element::Type& index_element_type)
+op::v10::Unique::Unique(const Output<Node>& data,
+                        const bool sorted,
+                        const element::Type& index_element_type,
+                        const element::Type& count_element_type)
     : op::Op{{data}},
       m_sorted{sorted},
-      m_index_element_type{index_element_type} {
+      m_index_element_type{index_element_type},
+      m_count_element_type{count_element_type} {
     constructor_validate_and_infer_types();
 }
 
 op::v10::Unique::Unique(const Output<Node>& data,
                         const Output<Node>& axis,
                         const bool sorted,
-                        const element::Type& index_element_type)
+                        const element::Type& index_element_type,
+                        const element::Type& count_element_type)
     : op::Op{{data, axis}},
       m_sorted{sorted},
-      m_index_element_type{index_element_type} {
+      m_index_element_type{index_element_type},
+      m_count_element_type{count_element_type} {
     constructor_validate_and_infer_types();
 }
 
@@ -114,6 +120,7 @@ bool op::v10::Unique::visit_attributes(AttributeVisitor& visitor) {
     OV_OP_SCOPE(v10_Unique_visit_attributes);
     visitor.on_attribute("sorted", m_sorted);
     visitor.on_attribute("index_element_type", m_index_element_type);
+    visitor.on_attribute("count_element_type", m_count_element_type);
     return true;
 }
 
@@ -121,7 +128,11 @@ void op::v10::Unique::validate_and_infer_types() {
     OV_OP_SCOPE(v10_Unique_validate_and_infer_types);
     NODE_VALIDATION_CHECK(this,
                           m_index_element_type == element::i32 || m_index_element_type == element::i64,
-                          "The element type of the outputs containing indices can only be set to i32 or i64");
+                          "The element type of the outputs containing indices can only be set to i32 or i64.");
+
+    NODE_VALIDATION_CHECK(this,
+                          m_count_element_type == element::i32 || m_count_element_type == element::i64,
+                          "The element type of the last output can only be set to i32 or i64.");
 
     const auto& input_shape = get_input_partial_shape(0);
     std::vector<PartialShape> output_shapes(4);
@@ -155,14 +166,15 @@ void op::v10::Unique::validate_and_infer_types() {
                 get_input_element_type(1) == element::i32 || get_input_element_type(1) == element::i64,
                 "The allowed element types of the 'axis' input tensor of the Unique operator are i32 and i64.");
 
-            NODE_VALIDATION_CHECK(
-                this,
-                get_input_partial_shape(1) == Shape{} || get_input_partial_shape(1) == Shape{1},
-                "The 'axis' input tensor of the Unique operator must be a scalar or 1D tensor with 1 element.");
-
             NODE_VALIDATION_CHECK(this,
                                   ov::op::util::is_constant(input_value(1).get_node()),
                                   "The 'axis' input of the Unique operator must be connected to a Constant.");
+
+            NODE_VALIDATION_CHECK(
+                this,
+                get_input_partial_shape(1) == PartialShape{} || get_input_partial_shape(1) == PartialShape{1},
+                "The 'axis' input tensor of the Unique operator must be a scalar or 1D tensor with 1 element.");
+
             const int64_t axis =
                 extract_axis(std::dynamic_pointer_cast<op::v0::Constant>(input_value(1).get_node_shared_ptr()));
 
@@ -211,19 +223,23 @@ void op::v10::Unique::validate_and_infer_types() {
     set_output_type(0, get_input_element_type(0), output_shapes[0]);
     set_output_type(1, m_index_element_type, output_shapes[1]);
     set_output_type(2, m_index_element_type, output_shapes[2]);
-    set_output_type(3, element::i64, output_shapes[3]);
+    set_output_type(3, m_count_element_type, output_shapes[3]);
 }
 
 std::shared_ptr<Node> op::v10::Unique::clone_with_new_inputs(const OutputVector& new_args) const {
     OV_OP_SCOPE(v10_Unique_clone_with_new_inputs);
     check_new_args_count(this, new_args);
     if (new_args.size() == 1) {
-        return std::make_shared<op::v10::Unique>(new_args.at(0), this->get_sorted(), this->get_index_element_type());
+        return std::make_shared<op::v10::Unique>(new_args.at(0),
+                                                 this->get_sorted(),
+                                                 this->get_index_element_type(),
+                                                 this->get_count_element_type());
     } else {
         return std::make_shared<op::v10::Unique>(new_args.at(0),
                                                  new_args.at(1),
                                                  this->get_sorted(),
-                                                 this->get_index_element_type());
+                                                 this->get_index_element_type(),
+                                                 this->get_count_element_type());
     }
 }
 }  // namespace ov

src/core/tests/type_prop/unique.cpp

@@ -48,6 +48,24 @@ TEST(type_prop, unique_no_axis_3d_index_type_i32) {
                         {{PartialShape{{1, 9}}, PartialShape{{1, 9}}, PartialShape{{9}}, PartialShape{{1, 9}}}});
 }
 
+TEST(type_prop, unique_no_axis_3d_count_type_i32) {
+    const auto data = make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 3});
+    const auto unique = make_shared<opset10::Unique>(data, true, element::i64, element::i32);
+
+    CHECK_ELEMENT_TYPES(unique, {{element::f32, element::i64, element::i64, element::i32}});
+    CHECK_OUTPUT_SHAPES(unique,
+                        {{PartialShape{{1, 9}}, PartialShape{{1, 9}}, PartialShape{{9}}, PartialShape{{1, 9}}}});
+}
+
+TEST(type_prop, unique_no_axis_3d_all_outputs_i32) {
+    const auto data = make_shared<opset10::Parameter>(element::f32, PartialShape{1, 3, 3});
+    const auto unique = make_shared<opset10::Unique>(data, true, element::i32, element::i32);
+
+    CHECK_ELEMENT_TYPES(unique, {{element::f32, element::i32, element::i32, element::i32}});
+    CHECK_OUTPUT_SHAPES(unique,
+                        {{PartialShape{{1, 9}}, PartialShape{{1, 9}}, PartialShape{{9}}, PartialShape{{1, 9}}}});
+}
+
 TEST(type_prop, unique_no_axis_scalar) {
     const auto data = make_shared<opset10::Parameter>(element::i32, PartialShape{});
     const auto unique = make_shared<opset10::Unique>(data);

src/core/tests/visitors/op/unique.cpp

@@ -19,9 +19,14 @@ TEST(attributes, unique_default_attributes) {
 
     const auto op = make_shared<opset10::Unique>(data);
     NodeBuilder builder(op);
+    auto g_unique = ov::as_type_ptr<opset10::Unique>(builder.create());
 
-    const auto expected_attr_count = 2;
+    constexpr auto expected_attr_count = 3;
     EXPECT_EQ(builder.get_value_map_size(), expected_attr_count);
+
+    EXPECT_EQ(op->get_sorted(), g_unique->get_sorted());
+    EXPECT_EQ(op->get_index_element_type(), g_unique->get_index_element_type());
+    EXPECT_EQ(op->get_count_element_type(), g_unique->get_count_element_type());
 }
 
 TEST(attributes, unique_sorted_false) {
@@ -33,7 +38,12 @@ TEST(attributes, unique_sorted_false) {
     NodeBuilder builder(op);
     auto g_unique = ov::as_type_ptr<opset10::Unique>(builder.create());
 
+    constexpr auto expected_attr_count = 3;
+    EXPECT_EQ(builder.get_value_map_size(), expected_attr_count);
+
     EXPECT_EQ(op->get_sorted(), g_unique->get_sorted());
+    EXPECT_EQ(op->get_index_element_type(), g_unique->get_index_element_type());
+    EXPECT_EQ(op->get_count_element_type(), g_unique->get_count_element_type());
 }
 
 TEST(attributes, unique_index_et_non_default) {
@@ -41,9 +51,14 @@ TEST(attributes, unique_index_et_non_default) {
     const auto data = make_shared<opset10::Parameter>(element::f32, Shape{1, 3, 10, 10});
     const auto grid = make_shared<opset10::Parameter>(element::f32, Shape{1, 5, 5, 2});
 
-    const auto op = make_shared<opset10::Unique>(data, true, element::i32);
+    const auto op = make_shared<opset10::Unique>(data, true, element::i32, element::i32);
     NodeBuilder builder(op);
     auto g_unique = ov::as_type_ptr<opset10::Unique>(builder.create());
 
+    constexpr auto expected_attr_count = 3;
+    EXPECT_EQ(builder.get_value_map_size(), expected_attr_count);
+
+    EXPECT_EQ(op->get_sorted(), g_unique->get_sorted());
     EXPECT_EQ(op->get_index_element_type(), g_unique->get_index_element_type());
+    EXPECT_EQ(op->get_count_element_type(), g_unique->get_count_element_type());
 }