Modify Tri class, revert tests

Author: Nimish-4

pytensor/tensor/basic.py

@@ -34,7 +34,7 @@
 from pytensor.npy_2_compat import normalize_axis_index, normalize_axis_tuple
 from pytensor.printing import Printer, min_informative_str, pprint, set_precedence
 from pytensor.raise_op import CheckAndRaise
-from pytensor.scalar import int32, upcast
+from pytensor.scalar import int32
 from pytensor.scalar.basic import ScalarConstant, ScalarType, ScalarVariable
 from pytensor.tensor import (
     _as_tensor_variable,
@@ -1148,14 +1148,6 @@ class Tri(OpFromGraph):
     Wrapper Op for np.tri graphs
     """
 
-    __props__ = ("dtype",)
-
-    def __init__(self, *args, M, k, dtype, **kwargs):
-        self.M = M
-        self.k = k
-        self.dtype = dtype
-        super().__init__(*args, **kwargs, strict=True)
-
 
 def tri(N, M=None, k=0, dtype=None):
     """
@@ -1184,14 +1176,19 @@ def tri(N, M=None, k=0, dtype=None):
     """
     if dtype is None:
         dtype = config.floatX
+    dtype = np.dtype(dtype)
+
     if M is None:
         M = N
 
+    N = as_tensor_variable(N)
+    M = as_tensor_variable(M)
+    k = as_tensor_variable(k)
+
     output = ((iota(as_tensor((N, 1)), 0) + k + 1) > iota(as_tensor((1, M)), 1)).astype(
         dtype
     )
-    N = as_tensor_variable(N)
-    return Tri(inputs=[N], outputs=[output], M=M, k=k, dtype=dtype)(N)
+    return Tri(inputs=[N, M, k], outputs=[output])(N, M, k)
 
 
 def tril(m, k=0):
@@ -1245,9 +1242,7 @@ def tril(m, k=0):
             [55, 56, 57, 58,  0]]])
 
     """
-    N, M = m.shape[-2:]
-    dtype = upcast(m.dtype)
-    return m * tri(N, M=M, k=k, dtype=dtype)  # M is symbolic, while it shouldnt be
+    return m * tri(*m.shape[-2:], k=k, dtype=m.dtype)
 
 
 def triu(m, k=0):

tests/tensor/test_basic.py

@@ -988,12 +988,15 @@ def check(dtype, N, M_=None, k=0):
             M = M_
             # Currently DebugMode does not support None as inputs even if this is
             # allowed.
-            if M is None and config.mode in ["DebugMode", "DEBUG_MODE"]:
+            if M is None:  # and config.mode in ["DebugMode", "DEBUG_MODE"]:
                 M = N
             N_symb = iscalar()
-            f = function([N_symb], tri(N_symb, M=M, k=k, dtype=dtype))
-            # kwargs = {}
-            result = f(N)
+            M_symb = iscalar()
+            k_symb = iscalar()
+            f = function(
+                [N_symb, M_symb, k_symb], tri(N_symb, M_symb, k_symb, dtype=dtype)
+            )
+            result = f(N, M, k)
             assert np.allclose(result, np.tri(N, M_, k, dtype=dtype))
             assert result.dtype == np.dtype(dtype)
 
@@ -1025,15 +1028,15 @@ def test_tril_triu(self):
 
         def check_l(m, k=0):
             m_symb = matrix(dtype=m.dtype)
-            # k_symb = iscalar()
-            f = function([m_symb], tril(m_symb, k=k))
+            k_symb = iscalar()
+            f = function([m_symb, k_symb], tril(m_symb, k_symb))
             f_indx = function(
-                [m_symb], tril_indices(m_symb.shape[0], k=k, m=m_symb.shape[1])
+                [m_symb, k_symb], tril_indices(m_symb.shape[0], k_symb, m_symb.shape[1])
             )
-            f_indx_from = function([m_symb], tril_indices_from(m_symb))
-            result = f(m)
-            result_indx = f_indx(m, k=k)
-            result_from = f_indx_from(m, k=k)
+            f_indx_from = function([m_symb, k_symb], tril_indices_from(m_symb, k_symb))
+            result = f(m, k)
+            result_indx = f_indx(m, k)
+            result_from = f_indx_from(m, k)
             assert np.allclose(result, np.tril(m, k))
             assert np.allclose(result_indx, np.tril_indices(m.shape[0], k, m.shape[1]))
             assert np.allclose(result_from, np.tril_indices_from(m, k))
@@ -1043,7 +1046,7 @@ def check_l(m, k=0):
         def check_u(m, k=0):
             m_symb = matrix(dtype=m.dtype)
             k_symb = iscalar()
-            f = function([m_symb, k_symb], triu(m_symb, k=k))
+            f = function([m_symb, k_symb], triu(m_symb, k_symb))
             f_indx = function(
                 [m_symb, k_symb], triu_indices(m_symb.shape[0], k_symb, m_symb.shape[1])
             )
@@ -1075,7 +1078,7 @@ def check_u_batch(m):
                 assert np.allclose(result, np.triu(m, k))
                 assert result.dtype == np.dtype(dtype)
 
-        for dtype in ["int32", "int64", "float32", "float64", "uint16", "complex64"]:
+        for dtype in ["int32", "int64", "float32", "float64", "uint16"]:
             m = random_of_dtype((10, 10), dtype)
             check_l(m, 0)
             check_l(m, 1)