Implement infer_shape method in Op class and remove redundant implementations in subclasses

Author: Aarsh-Wankar

pytensor/graph/op.py

@@ -20,6 +20,7 @@
     add_tag_trace,
     get_variable_trace_string,
 )
+from pytensor.tensor.utils import _gufunc_to_out_shape
 
 
 if TYPE_CHECKING:
@@ -596,6 +597,12 @@ def inplace_on_inputs(self, allowed_inplace_inputs: list[int]) -> "Op":
         # By default, do nothing
         return self
 
+    def infer_shape(self, fgraph, node, input_shapes):
+        if hasattr(self, "gufunc_signature"):
+            return _gufunc_to_out_shape(self.gufunc_signature, input_shapes)
+        else:
+            raise NotImplementedError(f"Op {self} does not implement infer_shape")
+
     def __str__(self):
         return getattr(type(self), "__name__", super().__str__())
 

pytensor/tensor/nlinalg.py

@@ -17,7 +17,6 @@
 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,9 +62,6 @@ def L_op(self, inputs, outputs, g_outputs):
         ).T
         return [grad]
 
-    def infer_shape(self, fgraph, node, shapes):
-        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
-
 
 def pinv(x, hermitian=False):
     """Computes the pseudo-inverse of a matrix :math:`A`.
@@ -156,9 +152,6 @@ def R_op(self, inputs, eval_points):
             return [None]
         return [-matrix_dot(xi, ev, xi)]
 
-    def infer_shape(self, fgraph, node, shapes):
-        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
-
 
 inv = matrix_inverse = Blockwise(MatrixInverse())
 
@@ -225,9 +218,6 @@ def grad(self, inputs, g_outputs):
         (x,) = inputs
         return [gz * self(x) * matrix_inverse(x).T]
 
-    def infer_shape(self, fgraph, node, shapes):
-        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
-
     def __str__(self):
         return "Det"
 
@@ -259,9 +249,6 @@ def perform(self, node, inputs, outputs):
         except Exception as e:
             raise ValueError("Failed to compute determinant", x) from e
 
-    def infer_shape(self, fgraph, node, shapes):
-        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
-
     def __str__(self):
         return "SLogDet"
 
@@ -317,9 +304,6 @@ def perform(self, node, inputs, outputs):
         (w, v) = outputs
         w[0], v[0] = (z.astype(x.dtype) for z in np.linalg.eig(x))
 
-    def infer_shape(self, fgraph, node, shapes):
-        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
-
 
 eig = Blockwise(Eig())
 

pytensor/tensor/slinalg.py

@@ -20,7 +20,6 @@
 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,9 +50,6 @@ def __init__(
         if self.overwrite_a:
             self.destroy_map = {0: [0]}
 
-    def infer_shape(self, fgraph, node, shapes):
-        return _gufunc_to_out_shape(self.gufunc_signature, shapes)
-
     def make_node(self, x):
         x = as_tensor_variable(x)
         if x.type.ndim != 2:
@@ -269,9 +265,6 @@ def make_node(self, A, b):
         x = tensor(dtype=o_dtype, shape=b.type.shape)
         return Apply(self, [A, b], [x])
 
-    def infer_shape(self, fgraph, node, shapes):
-        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`.
 
@@ -885,9 +878,6 @@ def perform(self, node, inputs, output_storage):
         out_dtype = node.outputs[0].type.dtype
         X[0] = scipy_linalg.solve_continuous_lyapunov(A, B).astype(out_dtype)
 
-    def infer_shape(self, fgraph, node, shapes):
-        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
         # Note that they write the equation as AX + XA.H + Q = 0, while scipy uses AX + XA^H = Q,
@@ -957,9 +947,6 @@ def perform(self, node, inputs, output_storage):
             out_dtype
         )
 
-    def infer_shape(self, fgraph, node, shapes):
-        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
         A, Q = inputs
@@ -1077,9 +1064,6 @@ def perform(self, node, inputs, output_storage):
         out_dtype = node.outputs[0].type.dtype
         X[0] = scipy_linalg.solve_discrete_are(A, B, Q, R).astype(out_dtype)
 
-    def infer_shape(self, fgraph, node, shapes):
-        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
         A, B, Q, R = inputs