Fix bug to ensure qml.math.dot works in autograph mode

doc/releases/changelog-dev.md

@@ -1026,6 +1026,9 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where `qml.math.dot` failed to work with `@tf.function` autograph mode.
+  [(#1842)](https://github.com/PennyLaneAI/pennylane/pull/1842)
+
 * Fixes a bug with the arrow width in the `measure` of `qml.circuit_drawer.MPLDrawer`. 
   [(#1823)](https://github.com/PennyLaneAI/pennylane/pull/1823)
 

pennylane/math/multi_dispatch.py

@@ -238,13 +238,13 @@ def dot(tensor1, tensor2):
         return np.tensordot(x, y, dims=[[-1], [-2]], like=interface)
 
     if interface == "tensorflow":
-        if x.ndim == 0 and y.ndim == 0:
+        if len(np.shape(x)) == 0 and len(np.shape(y)) == 0:
             return x * y
 
-        if y.ndim == 1:
+        if len(np.shape(y)) == 1:
             return np.tensordot(x, y, axes=[[-1], [0]], like=interface)
 
-        if x.ndim == 2 and y.ndim == 2:
+        if len(np.shape(x)) == 2 and len(np.shape(y)) == 2:
             return x @ y
 
         return np.tensordot(x, y, axes=[[-1], [-2]], like=interface)

tests/math/test_functions.py

@@ -475,6 +475,20 @@ def test_matrix_matrix_product(self, t1, t2):
         res = fn.dot(t1, t1)
         assert fn.allequal(res, np.array([[7, 10], [15, 22]]))
 
+    def test_matrix_vector_product_tensorflow_autograph(self):
+        """Test that the matrix-matrix dot product of two vectors results in a matrix
+        when using TensorFlow autograph mode"""
+        t1, t2 = tf.Variable([[1, 2], [3, 4]]), tf.Variable([6, 7])
+
+        @tf.function
+        def cost(t1, t2):
+            return fn.dot(t1, t2)
+
+        with tf.GradientTape() as tape:
+            res = cost(t1, t2)
+
+        assert fn.allequal(res, [20, 46])
+
     multidim_product_data = [
         [
             np.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),