simplify logic by making use of tensor.type.shape and removing broadcastable

Author: zoj613

aesara/tensor/nlinalg.py

@@ -836,45 +836,40 @@ def make_node(self, a, b):
 
         if 0 in (a.ndim, b.ndim):
             raise ValueError("inputs to `matmul` cannot be scalar.")
-        elif a.ndim >= 2 and b.ndim >= 2:
-            out = at.TensorType(
-                dtype=self.dtype,
-                broadcastable=(False,) * max(a.ndim, b.ndim),
-            )()
-        else:
-            out_ndim = b.ndim - 1 if a.ndim == 1 else a.ndim - 1
-            out = at.TensorType(
-                dtype=self.dtype,
-                broadcastable=(False,) * out_ndim,
-            )()
+
+        out_shape = self._get_output_shape(a, b, (a.type.shape, b.type.shape))
+        out = at.TensorType(dtype=self.dtype, shape=out_shape)()
         return Apply(self, [a, b], [out])
 
     def perform(self, node, inputs, outputs):
         x1, x2 = inputs
         outputs[0][0] = np.matmul(x1, x2, dtype=self.dtype)
 
-    def infer_shape(self, fgraph, node, shapes):
+    def _get_output_shape(self, x1, x2, shapes):
         x1_shape, x2_shape = shapes
-        x1, x2 = node.inputs
 
         if x1.ndim == 1 and x2.ndim == 1:
-            return [()]
+            return ()
         elif x1.ndim == 1 and x2.ndim > 1:
-            return [x2_shape[:-2] + x2_shape[-1:]]
+            return x2_shape[:-2] + x2_shape[-1:]
         elif x1.ndim > 1 and x2.ndim == 1:
-            return [x1_shape[:-1]]
+            return x1_shape[:-1]
         elif x1.ndim == 2 and x2.ndim == 2:
-            return [x1_shape[:-1] + x2_shape[-1:]]
+            return x1_shape[:-1] + x2_shape[-1:]
         elif x1.ndim > 2 and x2.ndim == 2:
-            return [x1_shape[:-2] + x1_shape[-2:-1] + x2_shape[-1:]]
+            return x1_shape[:-2] + x1_shape[-2:-1] + x2_shape[-1:]
         elif x1.ndim == 2 and x2.ndim > 2:
-            return [x2_shape[:-2] + x1_shape[-2:-1] + x2_shape[-1:]]
+            return x2_shape[:-2] + x1_shape[-2:-1] + x2_shape[-1:]
         else:
-            return [
+            return (
                 broadcast_shape(x1_shape[:-2], x2_shape[:-2], arrays_are_shapes=True)
                 + x1_shape[-2:-1]
                 + x2_shape[-1:]
-            ]
+            )
+
+    def infer_shape(self, fgraph, node, shapes):
+        x1, x2 = node.inputs
+        return [self._get_output_shape(x1, x2, shapes)]
 
 
 def matmul(x1, x2, dtype=None):

tests/tensor/test_nlinalg.py

@@ -183,8 +183,8 @@ def test_infer_shape(self):
             ((3, 2, 4, 5), (1, 5, 7)),
             ((4, 5, 8, 2, 3), (3, 2, 3, 4, 5, 8, 3, 3)),
         ]:
-            a = tensor(dtype=config.floatX, broadcastable=[i == 1 for i in shape_x1])
-            b = tensor(dtype=config.floatX, broadcastable=[i == 1 for i in shape_x2])
+            a = tensor(dtype=config.floatX, shape=shape_x1)
+            b = tensor(dtype=config.floatX, shape=shape_x2)
             x1 = self.rng.random(shape_x1).astype(config.floatX)
             x2 = self.rng.random(shape_x2).astype(config.floatX)