Fix for static PartialShape detection algorithm

ngraph/core/include/ngraph/partial_shape.hpp

@@ -230,11 +230,23 @@ namespace ngraph
         // True if the shape's rank is static.
         bool m_rank_is_static;
 
-        // True if the shape is static.
+        /// \brief Shape types. The shape type is lazily evaluated by calling the is_static()
+        /// method.
+        ///
+        /// \details It is highly recommended to avoid using the Dimension& operator[](size_t)
+        /// operator. It sets the shape type to SHAPE_IS_UPDATED and disables shape type caching.
+        /// Thus, the is_static method will have linear complexity because the shape is not
+        /// guaranteed to remain static or dynamic.
         mutable enum class ShapeType {
-            SHAPE_IS_UNKNOWN,
-            SHAPE_IS_STATIC,
-            SHAPE_IS_DYNAMIC
+            SHAPE_IS_UNKNOWN, // The shape type is unknown and should be calculated by checking all
+                              // dimensions.
+            SHAPE_IS_UPDATED, // User has retained a link to one dimension. Therefore, we can't
+                              // guarantee that the shape will remain static or dynamic, and its
+                              // type will always be evaluated.
+            SHAPE_IS_STATIC,  // The shape type is known and static. Also there are no any retained
+                              // dimensions by non-constant reference.
+            SHAPE_IS_DYNAMIC  // The shape type is dynamic and there are no any retained dimensions
+                              // by non-constant reference.
         } m_shape_type{ShapeType::SHAPE_IS_UNKNOWN};
 
         // Shape dimensions. This has no meaning if m_rank_is_static is false.

ngraph/core/src/partial_shape.cpp

@@ -60,16 +60,22 @@ PartialShape::PartialShape(const std::vector<Dimension>& dimensions)
 
 bool ngraph::PartialShape::is_static() const
 {
-    if (m_shape_type == ShapeType::SHAPE_IS_UNKNOWN)
+    ShapeType shape_type = m_shape_type;
+
+    if (m_shape_type == ShapeType::SHAPE_IS_UNKNOWN || m_shape_type == ShapeType::SHAPE_IS_UPDATED)
     {
-        m_shape_type =
+        shape_type =
             m_rank_is_static && std::all_of(m_dimensions.begin(),
                                             m_dimensions.end(),
                                             [](const Dimension& d) { return d.is_static(); })
                 ? ShapeType::SHAPE_IS_STATIC
                 : ShapeType::SHAPE_IS_DYNAMIC;
+
+        if (m_shape_type == ShapeType::SHAPE_IS_UNKNOWN)
+            m_shape_type = shape_type;
     }
-    return m_shape_type == ShapeType::SHAPE_IS_STATIC;
+
+    return shape_type == ShapeType::SHAPE_IS_STATIC;
 }
 
 bool ngraph::PartialShape::operator==(const PartialShape& partial_shape) const
@@ -473,7 +479,7 @@ Dimension& PartialShape::operator[](size_t i)
         throw std::out_of_range("Accessing out-of-range dimension in Dimension[]");
     }
     m_shape_type =
-        ShapeType::SHAPE_IS_UNKNOWN; // We can't guarantee that the shape remains static or dynamic.
+        ShapeType::SHAPE_IS_UPDATED; // We can't guarantee that the shape remains static or dynamic.
     return m_dimensions[i];
 }
 

ngraph/test/partial_shape.cpp

@@ -858,6 +858,23 @@ TEST(partial_shape, merge_rank_static_static_fail)
     ASSERT_TRUE(s.same_scheme(PartialShape{2, 3, Dimension::dynamic(), 5}));
 }
 
+TEST(partial_shape, changed_dimension_by_reference)
+{
+    PartialShape s{1, 2, 3};
+
+    Dimension& d = s[1];
+
+    ASSERT_TRUE(s.is_static());
+
+    d = Dimension::dynamic();
+
+    ASSERT_TRUE(s.is_dynamic());
+
+    d = 2;
+
+    ASSERT_TRUE(s.is_static());
+}
+
 TEST(partial_shape, infer_windowed_reduction_rank_dynamic_rank_dynamic_ok)
 {
     auto node = std::make_shared<op::Parameter>(element::f32, Shape{});