paramterized tests and added batch case

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -49,6 +49,7 @@
 
 
 logger = logging.getLogger(__name__)
+ALL_INVERSE_OPS = (MatrixInverse, MatrixPinv)
 
 
 def is_matrix_transpose(x: TensorVariable) -> bool:
@@ -593,11 +594,11 @@ def rewrite_inv_inv(fgraph, node):
     list of Variable, optional
         List of optimized variables, or None if no optimization was performed
     """
-    valid_inverses = (MatrixInverse, MatrixPinv)
+    ALL_INVERSE_OPS = (MatrixInverse, MatrixPinv)
     # Check if its a valid inverse operation (either inv/pinv)
     # In case the outer operation is an inverse, it directly goes to the next step of finding inner operation
     # If the outer operation is not a valid inverse, we do not apply this rewrite
-    if not isinstance(node.op.core_op, valid_inverses):
+    if not isinstance(node.op.core_op, ALL_INVERSE_OPS):
         return None
 
     potential_inner_inv = node.inputs[0].owner
@@ -608,7 +609,7 @@ def rewrite_inv_inv(fgraph, node):
     if not (
         potential_inner_inv
         and isinstance(potential_inner_inv.op, Blockwise)
-        and isinstance(potential_inner_inv.op.core_op, valid_inverses)
+        and isinstance(potential_inner_inv.op.core_op, ALL_INVERSE_OPS)
     ):
         return None
     return [potential_inner_inv.inputs[0]]
@@ -632,20 +633,19 @@ def rewrite_inv_eye_to_eye(fgraph, node):
     list of Variable, optional
         List of optimized variables, or None if no optimization was performed
     """
-    valid_inverses = (MatrixInverse, MatrixPinv)
     core_op = node.op.core_op
-    if not (isinstance(core_op, valid_inverses)):
+    if not (isinstance(core_op, ALL_INVERSE_OPS)):
         return None
 
     # Check whether input to inverse is Eye and the 1's are on main diagonal
-    eye_check = node.inputs[0]
+    potential_eye = node.inputs[0]
     if not (
-        eye_check.owner
-        and isinstance(eye_check.owner.op, Eye)
-        and getattr(eye_check.owner.inputs[-1], "data", -1).item() == 0
+        potential_eye.owner
+        and isinstance(potential_eye.owner.op, Eye)
+        and getattr(potential_eye.owner.inputs[-1], "data", -1).item() == 0
     ):
         return None
-    return [eye_check]
+    return [potential_eye]
 
 
 @register_canonicalize
@@ -668,9 +668,8 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
     list of Variable, optional
         List of optimized variables, or None if no optimization was performed
     """
-    valid_inverses = (MatrixInverse, MatrixPinv)
     core_op = node.op.core_op
-    if not (isinstance(core_op, valid_inverses)):
+    if not (isinstance(core_op, ALL_INVERSE_OPS)):
         return None
 
     inputs = node.inputs[0]
@@ -681,9 +680,8 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
         and AllocDiag.is_offset_zero(inputs.owner)
     ):
         inv_input = inputs.owner.inputs[0]
-        if inv_input.type.ndim == 1:
-            inv_val = pt.diag(1 / inv_input)
-            return [inv_val]
+        inv_val = pt.diag(1 / inv_input)
+        return [inv_val]
 
     # Check if the input is an elemwise multiply with identity matrix -- this also results in a diagonal matrix
     inputs_or_none = _find_diag_from_eye_mul(inputs)
@@ -700,8 +698,7 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
 
     # For a matrix, we have to first extract the diagonal (non-zero values) and then only use those
     if non_eye_input.type.broadcastable[-2:] == (False, False):
-        # For Matrix
-        return [eye_input / non_eye_input.diagonal(axis1=-1, axis2=-2)]
-    else:
-        # For Vector or Scalar
-        return [eye_input / non_eye_input]
+        non_eye_diag = non_eye_input.diagonal(axis1=-1, axis2=-2)
+        non_eye_input = pt.shape_padaxis(non_eye_diag, -2)
+
+    return [eye_input / non_eye_input]

tests/tensor/rewriting/test_linalg.py

@@ -560,22 +560,24 @@ def test_svd_uv_merge():
     assert svd_counter == 1
 
 
+def get_pt_function(x, op_name):
+    return getattr(pt.linalg, op_name)(x)
+
+
 @pytest.mark.parametrize("inv_op_1", ["inv", "pinv"])
 @pytest.mark.parametrize("inv_op_2", ["inv", "pinv"])
 def test_inv_inv_rewrite(inv_op_1, inv_op_2):
-    def get_pt_function(x, op_name):
-        return getattr(pt.linalg, op_name)(x)
-
     x = pt.matrix("x")
     op1 = get_pt_function(x, inv_op_1)
     op2 = get_pt_function(op1, inv_op_2)
     rewritten_out = rewrite_graph(op2)
     assert rewritten_out == x
 
 
-def test_inv_eye_to_eye():
+@pytest.mark.parametrize("inv_op", ["inv", "pinv"])
+def test_inv_eye_to_eye(inv_op):
     x = pt.eye(10)
-    x_inv = pt.linalg.inv(x)
+    x_inv = get_pt_function(x, inv_op)
     f_rewritten = function([], x_inv, mode="FAST_RUN")
     nodes = f_rewritten.maker.fgraph.apply_nodes
 
@@ -598,15 +600,16 @@ def test_inv_eye_to_eye():
 
 @pytest.mark.parametrize(
     "shape",
-    [(), (7,), (7, 7)],
-    ids=["scalar", "vector", "matrix"],
+    [(), (7,), (7, 7), (5, 7, 7)],
+    ids=["scalar", "vector", "matrix", "batched"],
 )
-def test_inv_diag_from_eye_mul(shape):
+@pytest.mark.parametrize("inv_op", ["inv", "pinv"])
+def test_inv_diag_from_eye_mul(shape, inv_op):
     # Initializing x based on scalar/vector/matrix
     x = pt.tensor("x", shape=shape)
     x_diag = pt.eye(7) * x
     # Calculating inverse using pt.linalg.inv
-    x_inv = pt.linalg.inv(x_diag)
+    x_inv = get_pt_function(x_diag, inv_op)
 
     # REWRITE TEST
     f_rewritten = function([x], x_inv, mode="FAST_RUN")
@@ -634,10 +637,11 @@ def test_inv_diag_from_eye_mul(shape):
     )
 
 
-def test_inv_diag_from_diag():
+@pytest.mark.parametrize("inv_op", ["inv", "pinv"])
+def test_inv_diag_from_diag(inv_op):
     x = pt.dvector("x")
     x_diag = pt.diag(x)
-    x_inv = pt.linalg.inv(x_diag)
+    x_inv = get_pt_function(x_diag, inv_op)
 
     # REWRITE TEST
     f_rewritten = function([x], x_inv, mode="FAST_RUN")