Allow single integer as TensorType shape

Author: ricardoV94

pytensor/tensor/type.py

@@ -71,7 +71,7 @@ class TensorType(CType[np.ndarray], HasDataType, HasShape):
     def __init__(
         self,
         dtype: str | npt.DTypeLike,
-        shape: Iterable[bool | int | None] | None = None,
+        shape: Iterable[bool | int | None] | int | None = None,
         name: str | None = None,
         broadcastable: Iterable[bool] | None = None,
     ):
@@ -99,7 +99,7 @@ def __init__(
             )
             shape = broadcastable
 
-        if str(dtype) == "floatX":
+        if dtype == "floatX":
             self.dtype = config.floatX
         else:
             try:
@@ -118,6 +118,8 @@ def parse_bcast_and_shape(s):
                 f"TensorType broadcastable/shape must be a boolean, integer or None, got {type(s)} {s}"
             )
 
+        if isinstance(shape, int):
+            shape = (shape,)
         self.shape = tuple(parse_bcast_and_shape(s) for s in shape)
         self.dtype_specs()  # error checking is done there
         self.name = name