Allow subtypes by default in make_node

aesara/graph/op.py

@@ -222,21 +222,22 @@ class Op(MetaObject):
         """
         if self.itypes is None:
             raise NotImplementedError(
-                "You can either define itypes and otypes,\
-             or implement make_node"
+                "You can either define itypes and otypes, or implement make_node"
             )
 
         if self.otypes is None:
             raise NotImplementedError(
-                "You can either define itypes and otypes,\
-             or implement make_node"
+                "You can either define itypes and otypes, or implement make_node"
             )
 
         if len(inputs) != len(self.itypes):
             raise ValueError(
                 f"We expected {len(self.itypes)} inputs but got {len(inputs)}."
             )
-        if not all(it.in_same_class(inp.type) for inp, it in zip(inputs, self.itypes)):
+        if not all(
+            expected_type.is_super(var.type)
+            for var, expected_type in zip(inputs, self.itypes)
+        ):
             raise TypeError(
                 f"Invalid input types for Op {self}:\n"
                 + "\n".join(

tests/graph/test_op.py

@@ -223,3 +223,16 @@ def test_op_invalid_input_types():
     msg = r"^Invalid input types for Op.*"
     with pytest.raises(TypeError, match=msg):
         TestOp()(dvector(), dscalar(), dvector())
+
+
+def test_op_input_broadcastable():
+    # Test that we can create an op with a broadcastable subtype as input
+    class SomeOp(aesara.tensor.Op):
+        itypes = [at.dvector]
+        otypes = [at.dvector]
+
+        def perform(self, *_):
+            raise NotImplementedError()
+
+    x = at.TensorType(dtype="float64", shape=(1,))("x")
+    assert SomeOp()(x).type == at.dvector