Fixed SolveTriangular Op in PyTorch

Added support for a b param of ndim = 1.

Author: twaclaw

Diff for: pytensor/link/pytorch/dispatch/slinalg.py

@@ -71,10 +71,19 @@ def solve_triangular(A, b):
         if trans == 2 or trans == "C":
             A_p = A.H
 
-        return torch.linalg.solve_triangular(
-            A_p, b, upper=upper, unitriangular=unit_diagonal
+        b_p = b
+        if b.ndim == 1:
+            b_p = b[:, None]
+
+        res = torch.linalg.solve_triangular(
+            A_p, b_p, upper=upper, unitriangular=unit_diagonal
         )
 
+        if b.ndim == 1 and res.shape[1] == 1:
+            return res.flatten()
+
+        return res
+
     return solve_triangular
 
 

Diff for: tests/link/pytorch/test_slinalg.py

@@ -57,7 +57,7 @@ def test_pytorch_basic():
         ],
     )
 
-    out = pt_slinalg.Cholesky(lower=False)(x)
+    out = pt_slinalg.cholesky(x, lower=False)
     out_fg = FunctionGraph([x], [out])
     compare_pytorch_and_py(
         out_fg,
@@ -79,7 +79,6 @@ def test_pytorch_basic():
     )
 
 
-@pytest.mark.xfail(reason="Blockwise not implemented")
 @pytest.mark.parametrize(
     "check_finite",
     (False, pytest.param(True, marks=pytest.mark.xfail(raises=NotImplementedError))),