updated tests

Author: tanish1729

tests/tensor/rewriting/test_linalg.py

@@ -41,6 +41,9 @@
 from tests.test_rop import break_op
 
 
+ATOL = RTOL = 1e-3 if config.floatX == "float32" else 1e-8
+
+
 def test_rop_lop():
     mx = matrix("mx")
     mv = matrix("mv")
@@ -568,3 +571,90 @@ def get_pt_function(x, op_name):
     op2 = get_pt_function(op1, inv_op_2)
     rewritten_out = rewrite_graph(op2)
     assert rewritten_out == x
+
+
+def test_inv_eye_to_eye():
+    x = pt.eye(10)
+    x_inv = pt.linalg.inv(x)
+    f_rewritten = function([], x_inv, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+
+    # Rewrite Test
+    valid_inverses = (MatrixInverse, MatrixPinv)
+    assert not any(isinstance(node.op, valid_inverses) for node in nodes)
+
+    # Value Test
+    x_test = np.eye(10)
+    x_inv_val = np.linalg.inv(x_test)
+    rewritten_val = f_rewritten()
+
+    assert_allclose(
+        x_inv_val,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [(), (7,), (7, 7)],
+    ids=["scalar", "vector", "matrix"],
+)
+def test_inv_diag_from_eye_mul(shape):
+    # 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)
+
+    # REWRITE TEST
+    f_rewritten = function([x], x_inv, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+
+    valid_inverses = (MatrixInverse, MatrixPinv)
+    assert not any(isinstance(node.op, valid_inverses) for node in nodes)
+
+    # NUMERIC VALUE TEST
+    if len(shape) == 0:
+        x_test = np.array(np.random.rand()).astype(config.floatX)
+    elif len(shape) == 1:
+        x_test = np.random.rand(*shape).astype(config.floatX)
+    else:
+        x_test = np.random.rand(*shape).astype(config.floatX)
+    x_test_matrix = np.eye(7) * x_test
+    inverse_matrix = np.linalg.inv(x_test_matrix)
+    rewritten_inverse = f_rewritten(x_test)
+
+    assert_allclose(
+        inverse_matrix,
+        rewritten_inverse,
+        atol=ATOL,
+        rtol=RTOL,
+    )
+
+
+def test_inv_diag_from_diag():
+    x = pt.dvector("x")
+    x_diag = pt.diag(x)
+    x_inv = pt.linalg.inv(x_diag)
+
+    # REWRITE TEST
+    f_rewritten = function([x], x_inv, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+
+    valid_inverses = (MatrixInverse, MatrixPinv)
+    assert not any(isinstance(node.op, valid_inverses) for node in nodes)
+
+    # NUMERIC VALUE TEST
+    x_test = np.random.rand(10)
+    x_test_matrix = np.eye(10) * x_test
+    inverse_matrix = np.linalg.inv(x_test_matrix)
+    rewritten_inverse = f_rewritten(x_test)
+
+    assert_allclose(
+        inverse_matrix,
+        rewritten_inverse,
+        atol=ATOL,
+        rtol=RTOL,
+    )