Create is_unknown and has_unknown_dim on Shape (#201)

This PR enables usages described in
#199.

This pull request enhances the handling and documentation of unknown
dimensions in the `Shape` class within the ONNX IR core. It introduces
new methods to check for unknown dimensions, updates the documentation
to clarify shape equality with unknowns, and adds comprehensive tests to
ensure correct behavior. Additionally, it clarifies the serialization
logic for symbolic dimensions with unknown values.

**Shape API improvements:**

* Added `is_unknown(dim)` and `has_unknown_dim()` methods to the `Shape`
class to check for unknown (i.e., `None`) dimensions, and updated the
class docstring to clarify shape equality semantics when unknowns are
present.

**Serialization:**

* Clarified that a symbolic dimension with `None` value is valid and
should result in an empty field in the serialized proto, reflecting ONNX
semantics for unknown dimensions.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

src/onnx_ir/_core.py

@@ -1219,6 +1219,12 @@ class Shape(_protocols.ShapeProtocol, _display.PrettyPrintable):
 
     Use :meth:`get_denotation` and :meth:`set_denotation` to access and modify the denotations.
 
+    .. note::
+        Two shapes can be compared for equality. Be careful when comparing shapes with
+        unknown dimensions (``None``), as they may not be considered semantically equal
+        even if all dimensions are the same. You can use :meth:`has_unknown_dim` to
+        check if a shape has any unknown dimensions.
+
     Example::
 
         >>> import onnx_ir as ir
@@ -1427,6 +1433,29 @@ def is_dynamic(self, dim=None) -> bool:
             return not self.is_static()
         return not self.is_static(dim)
 
+    def is_unknown_dim(self, dim: int) -> bool:
+        """Return True if the dimension is unknown (None).
+
+        A dynamic dimension without a symbolic name is considered unknown.
+
+        .. versionadded:: 0.1.10
+
+        Args:
+            dim: The index of the dimension.
+        """
+        dim_obj = self._dims[dim]
+        return isinstance(dim_obj, SymbolicDim) and dim_obj.value is None
+
+    def has_unknown_dim(self) -> bool:
+        """Return True if any dimension is unknown (None).
+
+        You can use :meth:`is_unknown_dim` to check if a specific dimension is unknown.
+
+        .. versionadded:: 0.1.10
+        """
+        # We can use "in" directly because SymbolicDim implements __eq__ with None
+        return None in self._dims
+
 
 def _quoted(string: str) -> str:
     """Return a quoted string.

src/onnx_ir/_core_test.py

@@ -743,6 +743,45 @@ def test_is_dynamic_on_empty_shape(self):
         shape = _core.Shape(())
         self.assertFalse(shape.is_dynamic())
 
+    def test_is_unknown_dim(self):
+        shape = _core.Shape([42, None, "any string", None])
+        self.assertFalse(shape.is_unknown_dim(0))  # integer dimension is not unknown
+        self.assertTrue(shape.is_unknown_dim(1))  # None dimension is unknown
+        self.assertFalse(
+            shape.is_unknown_dim(2)
+        )  # string dimension is not unknown (it's symbolic)
+        self.assertTrue(shape.is_unknown_dim(3))  # None dimension is unknown
+
+    def test_is_unknown_dim_raises_when_index_out_of_range(self):
+        shape = _core.Shape([42])
+        with self.assertRaises(IndexError):
+            shape.is_unknown_dim(1)
+
+    def test_has_unknown_dim(self):
+        # Shape with unknown dimensions
+        shape = _core.Shape([42, None, "any string"])
+        self.assertTrue(shape.has_unknown_dim())
+
+        # Shape with only None dimensions
+        shape = _core.Shape([None, None])
+        self.assertTrue(shape.has_unknown_dim())
+
+        # Shape with no unknown dimensions (static and symbolic)
+        shape = _core.Shape([42, "any string", 64])
+        self.assertFalse(shape.has_unknown_dim())
+
+        # Shape with only static dimensions
+        shape = _core.Shape([42, 64, 128])
+        self.assertFalse(shape.has_unknown_dim())
+
+        # Shape with only symbolic dimensions
+        shape = _core.Shape(["batch", "height", "width"])
+        self.assertFalse(shape.has_unknown_dim())
+
+    def test_has_unknown_dim_on_empty_shape(self):
+        shape = _core.Shape(())
+        self.assertFalse(shape.has_unknown_dim())
+
 
 class ValueTest(unittest.TestCase):
     def setUp(self) -> None:

src/onnx_ir/serde.py

@@ -2000,5 +2000,8 @@ def serialize_dimension_into(
         dim_proto.dim_value = dim
     elif isinstance(dim, (_core.SymbolicDim, _protocols.SymbolicDimProtocol)):
         if dim.value is not None:
-            # TODO(justinchuby): None is probably not a valid value for dim_param
             dim_proto.dim_param = str(dim.value)
+        # NOTE: None is a valid value for symbolic dimension:
+        # A dimension MAY have neither dim_value nor dim_param set. Such a dimension
+        # represents an unknown dimension unrelated to other unknown dimensions.
+        # Here we will just leave the dim_proto empty.