Do not call tensordot when general_dot is just an elemwise multiplication

Author: ricardoV94

pytensor/tensor/einsum.py

@@ -131,15 +131,13 @@ def _general_dot(
     # TODO: tensordot produces very complicated graphs unnecessarily
     #  In some cases we are just doing elemwise addition after some transpositions
     #  We also have some Blockwise(Reshape) that will slow down things!
-    out = vectorize(
-        partial(tensordot, axes=[core_lhs_axes, core_rhs_axes]), signature=signature
-    )(lhs, rhs)
-
-    # # Reorder batch axes according to the original order of lhs
-    # original_lhs_batch_axes, _ = batch_axes
-    # final_batch_axes = tuple(np.argsort(original_lhs_batch_axes))
-    # new_batch_axes = tuple(range(lhs_n_batch_axes))
-    # out = moveaxis(out, new_batch_axes, final_batch_axes)
+    if signature == "(),()->()":
+        # Just a multiplication
+        out = lhs * rhs
+    else:
+        out = vectorize(
+            partial(tensordot, axes=[core_lhs_axes, core_rhs_axes]), signature=signature
+        )(lhs, rhs)
 
     return cast(TensorVariable, out)