Modify tests for tri Op

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
+from pytensor.scalar import int32, upcast
 from pytensor.scalar.basic import ScalarConstant, ScalarType, ScalarVariable
 from pytensor.tensor import (
     _as_tensor_variable,
@@ -1186,8 +1186,9 @@ def tri(N, M=None, k=0, dtype=None):
         dtype = config.floatX
     if M is None:
         M = N
+
     output = ((iota(as_tensor((N, 1)), 0) + k + 1) > iota(as_tensor((1, M)), 1)).astype(
-        int
+        dtype
     )
     N = as_tensor_variable(N)
     return Tri(inputs=[N], outputs=[output], M=M, k=k, dtype=dtype)(N)
@@ -1244,7 +1245,9 @@ def tril(m, k=0):
             [55, 56, 57, 58,  0]]])
 
     """
-    return m * tri(*m.shape[-2:], k=k, dtype=m.dtype)
+    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
 
 
 def triu(m, k=0):

tests/tensor/test_basic.py

@@ -991,16 +991,19 @@ def check(dtype, N, M_=None, k=0):
             if M is None and config.mode in ["DebugMode", "DEBUG_MODE"]:
                 M = N
             N_symb = iscalar()
-            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)
+            f = function([N_symb], tri(N_symb, M=M, k=k, dtype=dtype))
+            # kwargs = {}
+            result = f(N)
             assert np.allclose(result, np.tri(N, M_, k, dtype=dtype))
             assert result.dtype == np.dtype(dtype)
 
-        for dtype in ["int32", "int64", "float32", "float64", "uint16", "complex64"]:
+        for dtype in [
+            "int32",
+            "int64",
+            "float32",
+            "float64",
+            "uint16",
+        ]:  # Handle "complex64" ?
             check(dtype, 3)
             # M != N, k = 0
             check(dtype, 3, 5)
@@ -1022,15 +1025,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, k_symb], tril(m_symb, k_symb))
+            # k_symb = iscalar()
+            f = function([m_symb], tril(m_symb, k=k))
             f_indx = function(
-                [m_symb, k_symb], tril_indices(m_symb.shape[0], k_symb, m_symb.shape[1])
+                [m_symb], tril_indices(m_symb.shape[0], k=k, m=m_symb.shape[1])
             )
-            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)
+            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)
             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))
@@ -1040,7 +1043,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_symb))
+            f = function([m_symb, k_symb], triu(m_symb, k=k))
             f_indx = function(
                 [m_symb, k_symb], triu_indices(m_symb.shape[0], k_symb, m_symb.shape[1])
             )