Reverted folder structure and added BatchedDot

Author: HangenYuu

pytensor/link/pytorch/dispatch/__init__.py

@@ -2,6 +2,9 @@
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify, pytorch_typify
 
 # # Load dispatch specializations
+import pytensor.link.pytorch.dispatch.blas
 import pytensor.link.pytorch.dispatch.scalar
 import pytensor.link.pytorch.dispatch.elemwise
+import pytensor.link.pytorch.dispatch.math
+
 # isort: on

pytensor/link/pytorch/dispatch/blas.py

@@ -0,0 +1,14 @@
+import torch
+
+from pytensor.link.pytorch.dispatch import pytorch_funcify
+from pytensor.tensor.blas import BatchedDot
+
+
+@pytorch_funcify.register(BatchedDot)
+def pytorch_funcify_BatchedDot(op, **kwargs):
+    def batched_dot(a, b):
+        if a.shape[0] != b.shape[0]:
+            raise TypeError("Shapes must match in the 0-th dimension")
+        return torch.bmm(a, b)
+
+    return batched_dot

pytensor/link/pytorch/dispatch/nlinalg.py renamed to pytensor/link/pytorch/dispatch/math.py

@@ -0,0 +1,36 @@
+import numpy as np
+import pytest
+
+from pytensor.compile.function import function
+from pytensor.compile.mode import Mode
+from pytensor.configdefaults import config
+from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.op import get_test_value
+from pytensor.graph.rewriting.db import RewriteDatabaseQuery
+from pytensor.link.pytorch import PytorchLinker
+from pytensor.tensor import blas as pt_blas
+from pytensor.tensor.type import tensor3
+from tests.link.pytorch.test_basic import compare_pytorch_and_py
+
+
+def test_pytorch_BatchedDot():
+    # tensor3 . tensor3
+    a = tensor3("a")
+    a.tag.test_value = (
+        np.linspace(-1, 1, 10 * 5 * 3).astype(config.floatX).reshape((10, 5, 3))
+    )
+    b = tensor3("b")
+    b.tag.test_value = (
+        np.linspace(1, -1, 10 * 3 * 2).astype(config.floatX).reshape((10, 3, 2))
+    )
+    out = pt_blas.BatchedDot()(a, b)
+    fgraph = FunctionGraph([a, b], [out])
+    compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+
+    # A dimension mismatch should raise a TypeError for compatibility
+    inputs = [get_test_value(a)[:-1], get_test_value(b)]
+    opts = RewriteDatabaseQuery(include=[None], exclude=["cxx_only", "BlasOpt"])
+    pytorch_mode = Mode(PytorchLinker(), opts)
+    pytensor_jax_fn = function(fgraph.inputs, fgraph.outputs, mode=pytorch_mode)
+    with pytest.raises(TypeError):
+        pytensor_jax_fn(*inputs)