Changed implementation of dot. Renamed tests

Author: HangenYuu

pytensor/link/pytorch/dispatch/nlinalg.py

@@ -7,6 +7,11 @@
 @pytorch_funcify.register(Dot)
 def pytorch_funcify_Dot(op, **kwargs):
     def dot(x, y):
-        return torch.matmul(x, y)
+        # Case 1: Vector Product/Matrix Multiplication/1-D Broadcastable Vector
+        if x.shape < 3 and y.shape < 3:
+            return torch.matmul(x, y)
+        else:
+            # Case 2: Stackable batch dimension
+            return torch.tensordot(x, y, dims=([-1], [-2]))
 
     return dot

tests/link/pytorch/test_nlinalg.py

@@ -3,11 +3,15 @@
 from pytensor.configdefaults import config
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import get_test_value
-from pytensor.tensor.type import matrix, scalar, vector
+from pytensor.tensor.type import matrix, scalar, tensor3, vector
 from tests.link.pytorch.test_basic import compare_pytorch_and_py
 
 
-def test_tensor_basics():
+def test_pytorch_dot():
+    a = tensor3("a")
+    a.tag.test_value = np.zeros((3, 2, 4)).astype(config.floatX)
+    b = tensor3("b")
+    b.tag.test_value = np.zeros((3, 4, 1)).astype(config.floatX)
     y = vector("y")
     y.tag.test_value = np.r_[1.0, 2.0].astype(config.floatX)
     x = vector("x")
@@ -19,12 +23,17 @@ def test_tensor_basics():
     beta = scalar("beta")
     beta.tag.test_value = np.array(5.0, dtype=config.floatX)
 
-    # 1D * 2D * 1D
-    out = y.dot(alpha * A).dot(x) + beta * y
-    fgraph = FunctionGraph([y, x, A, alpha, beta], [out])
+    # 3D * 3D
+    out = a.dot(b * alpha) + beta * b
+    fgraph = FunctionGraph([a, b, alpha, beta], [out])
     compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
 
     # 2D * 2D
     out = A.dot(A * alpha) + beta * A
     fgraph = FunctionGraph([A, alpha, beta], [out])
     compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+
+    # 1D * 2D and 1D * 1D
+    out = y.dot(alpha * A).dot(x) + beta * y
+    fgraph = FunctionGraph([y, x, A, alpha, beta], [out])
+    compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])