Fix triu docstring; remove dtype prop in JAX version

Nimish-4 · Nimish-4 · commit 00ab6b246632 · 2025-03-14T20:23:50.000+05:30
diff --git a/pytensor/link/jax/dispatch/tensor_basic.py b/pytensor/link/jax/dispatch/tensor_basic.py
@@ -203,6 +203,6 @@ def tri(*args):
             x if const_x is None else const_x
             for x, const_x in zip(args, const_args, strict=True)
         ]
-        return jnp.tri(*args, dtype=op.dtype)
+        return jnp.tri(*args)
 
     return tri
diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py
@@ -1268,7 +1268,7 @@ def triu(m, k=0):
            [ 0,  8,  9],
            [ 0,  0, 12]])
 
-    >>> pt.triu(np.arange(3 * 4 * 5).reshape((3, 4, 5))).eval()
+    >>> pt.triu(pt.arange(3 * 4 * 5).reshape((3, 4, 5))).eval()
     array([[[ 0,  1,  2,  3,  4],
             [ 0,  6,  7,  8,  9],
             [ 0,  0, 12, 13, 14],

Original file line number	Diff line number	Diff line change
`@@ -203,6 +203,6 @@ def tri(*args):`
`203`	`203`	`x if const_x is None else const_x`
`204`	`204`	`for x, const_x in zip(args, const_args, strict=True)`
`205`	`205`	`]`
`206`		`- return jnp.tri(*args, dtype=op.dtype)`
	`206`	`+ return jnp.tri(*args)`
`207`	`207`
`208`	`208`	`return tri`