Refactor infer_shape methods to utilize _gufunc_to_out_shape for output shape computation

Author: Aarsh-Wankar

pytensor/tensor/nlinalg.py

@@ -17,6 +17,7 @@
 from pytensor.tensor.basic import as_tensor_variable, diagonal
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.type import Variable, dvector, lscalar, matrix, scalar, vector
+from pytensor.tensor.utils import _gufunc_to_out_shape
 
 
 class MatrixPinv(Op):
@@ -63,7 +64,7 @@ def L_op(self, inputs, outputs, g_outputs):
         return [grad]
 
     def infer_shape(self, fgraph, node, shapes):
-        return [list(reversed(shapes[0]))]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
 
 def pinv(x, hermitian=False):
@@ -156,7 +157,7 @@ def R_op(self, inputs, eval_points):
         return [-matrix_dot(xi, ev, xi)]
 
     def infer_shape(self, fgraph, node, shapes):
-        return shapes
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
 
 inv = matrix_inverse = Blockwise(MatrixInverse())
@@ -225,7 +226,7 @@ def grad(self, inputs, g_outputs):
         return [gz * self(x) * matrix_inverse(x).T]
 
     def infer_shape(self, fgraph, node, shapes):
-        return [()]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def __str__(self):
         return "Det"
@@ -259,7 +260,7 @@ def perform(self, node, inputs, outputs):
             raise ValueError("Failed to compute determinant", x) from e
 
     def infer_shape(self, fgraph, node, shapes):
-        return [(), ()]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def __str__(self):
         return "SLogDet"
@@ -317,8 +318,7 @@ def perform(self, node, inputs, outputs):
         w[0], v[0] = (z.astype(x.dtype) for z in np.linalg.eig(x))
 
     def infer_shape(self, fgraph, node, shapes):
-        n = shapes[0][0]
-        return [(n,), (n, n)]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
 
 eig = Blockwise(Eig())
@@ -619,16 +619,7 @@ def perform(self, node, inputs, outputs):
             s[0] = np.linalg.svd(x, self.full_matrices, self.compute_uv)
 
     def infer_shape(self, fgraph, node, shapes):
-        (x_shape,) = shapes
-        M, N = x_shape
-        K = ptm.minimum(M, N)
-        s_shape = (K,)
-        if self.compute_uv:
-            u_shape = (M, M) if self.full_matrices else (M, K)
-            vt_shape = (N, N) if self.full_matrices else (K, N)
-            return [u_shape, s_shape, vt_shape]
-        else:
-            return [s_shape]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def L_op(
         self,

pytensor/tensor/slinalg.py

@@ -20,6 +20,7 @@
 from pytensor.tensor.nlinalg import kron, matrix_dot
 from pytensor.tensor.shape import reshape
 from pytensor.tensor.type import matrix, tensor, vector
+from pytensor.tensor.utils import _gufunc_to_out_shape
 from pytensor.tensor.variable import TensorVariable
 
 
@@ -51,7 +52,7 @@ def __init__(
             self.destroy_map = {0: [0]}
 
     def infer_shape(self, fgraph, node, shapes):
-        return [shapes[0]]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def make_node(self, x):
         x = as_tensor_variable(x)
@@ -269,13 +270,7 @@ def make_node(self, A, b):
         return Apply(self, [A, b], [x])
 
     def infer_shape(self, fgraph, node, shapes):
-        Ashape, Bshape = shapes
-        rows = Ashape[1]
-        if len(Bshape) == 1:
-            return [(rows,)]
-        else:
-            cols = Bshape[1]
-            return [(rows, cols)]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def L_op(self, inputs, outputs, output_gradients):
         r"""Reverse-mode gradient updates for matrix solve operation :math:`c = A^{-1} b`.
@@ -891,7 +886,7 @@ def perform(self, node, inputs, output_storage):
         X[0] = scipy_linalg.solve_continuous_lyapunov(A, B).astype(out_dtype)
 
     def infer_shape(self, fgraph, node, shapes):
-        return [shapes[0]]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def grad(self, inputs, output_grads):
         # Gradient computations come from Kao and Hennequin (2020), https://arxiv.org/pdf/2011.11430.pdf
@@ -963,7 +958,7 @@ def perform(self, node, inputs, output_storage):
         )
 
     def infer_shape(self, fgraph, node, shapes):
-        return [shapes[0]]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def grad(self, inputs, output_grads):
         # Gradient computations come from Kao and Hennequin (2020), https://arxiv.org/pdf/2011.11430.pdf
@@ -1083,7 +1078,7 @@ def perform(self, node, inputs, output_storage):
         X[0] = scipy_linalg.solve_discrete_are(A, B, Q, R).astype(out_dtype)
 
     def infer_shape(self, fgraph, node, shapes):
-        return [shapes[0]]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def grad(self, inputs, output_grads):
         # Gradient computations come from Kao and Hennequin (2020), https://arxiv.org/pdf/2011.11430.pdf
@@ -1181,8 +1176,7 @@ def grad(self, inputs, gout):
         return [gout[0][slc] for slc in slices]
 
     def infer_shape(self, fgraph, nodes, shapes):
-        first, second = zip(*shapes, strict=True)
-        return [(pt.add(*first), pt.add(*second))]
+        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
 
     def _validate_and_prepare_inputs(self, matrices, as_tensor_func):
         if len(matrices) != self.n_inputs:

pytensor/tensor/utils.py

@@ -202,6 +202,39 @@ def _parse_gufunc_signature(
     )
 
 
+def _gufunc_to_out_shape(
+    signature: str, shapes: list[tuple[int, ...]]
+) -> list[tuple[int, ...]]:
+    """
+    Compute the shape of the output of an Op given its gufunc signature and the
+    shapes of its inputs.
+
+    Parameters
+    ----------
+    signature : str
+        The gufunc signature of the Op.
+        eg: "(m,n),(n,p)->(m,p)".
+
+    shapes : list of tuple of int
+        The list of shapes of the inputs.
+
+    Returns
+    -------
+    out_shape : list of tuple of int
+        The list of shapes of the outputs.
+    """
+    parsed = _parse_gufunc_signature(signature)
+    out_shape = []
+    dic = dict()
+    for i in range(len(parsed[0])):
+        for j in range(len(parsed[0][i])):
+            dic[parsed[0][i][j]] = shapes[i][j]
+    for i in range(len(parsed[1])):
+        temp_list = [dic[x] for x in parsed[1][i]]
+        out_shape.append(tuple(temp_list))
+    return out_shape
+
+
 def safe_signature(
     core_inputs_ndim: Sequence[int],
     core_outputs_ndim: Sequence[int],