Add HasShape mixin

Author: brandonwillard

aesara/graph/type.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Any, Optional, Text, TypeVar, Union
+from typing import Any, Optional, Text, Tuple, TypeVar, Union
 
 from typing_extensions import TypeAlias
 
@@ -257,6 +257,13 @@ def values_eq_approx(cls, a: Any, b: Any):
 
 
 class HasDataType:
-    """A mixing for an `Op` type that has a :attr:`dtype` attribute."""
+    """A mixin for a type that has a :attr:`dtype` attribute."""
 
     dtype: str
+
+
+class HasShape:
+    """A mixin for a type that has :attr:`shape` and :attr:`ndim` attributes."""
+
+    ndim: int
+    shape: Tuple[Optional[int], ...]

aesara/tensor/type.py

@@ -8,7 +8,7 @@
 from aesara import scalar as aes
 from aesara.configdefaults import config
 from aesara.graph.basic import Variable
-from aesara.graph.type import HasDataType
+from aesara.graph.type import HasDataType, HasShape
 from aesara.graph.utils import MetaType
 from aesara.link.c.type import CType
 from aesara.misc.safe_asarray import _asarray
@@ -48,7 +48,7 @@
 }
 
 
-class TensorType(CType, HasDataType):
+class TensorType(CType, HasDataType, HasShape):
     r"""Symbolic `Type` representing `numpy.ndarray`\s."""
 
     __props__: Tuple[str, ...] = ("dtype", "shape")