Implement vectorize for xtensor Dot

Author: ricardoV94

pytensor/xtensor/math.py

@@ -566,6 +566,9 @@ def make_node(self, x, y):
         out = xtensor(dtype=out_dtype, shape=out_shape, dims=out_dims)
         return Apply(self, [x, y], [out])
 
+    def vectorize_node(self, node, *new_inputs, new_dim):
+        return self(*new_inputs, return_list=True)
+
 
 def dot(x, y, dim: str | Sequence[str] | EllipsisType | None = None):
     """Generalized dot product for XTensorVariables.

tests/xtensor/test_math.py

@@ -353,3 +353,13 @@ def test_xelemwise_vectorize():
 
     check_vectorization([ab], [exp(ab)])
     check_vectorization([ab, bc], [ab + bc])
+
+
+def test_dot_vectorize():
+    x = xtensor("x", dims=("a", "b"), shape=(2, 3))
+    y = xtensor("y", dims=("b", "c"), shape=(3, 4))
+
+    check_vectorization([x, y], [x.dot(y)])
+    check_vectorization([x, y], [x.dot(y, dim=("a", "b"))])
+    check_vectorization([x, y], [x.dot(y, dim="c")])
+    check_vectorization([x, y], [x.dot(y, dim=...)])