Repository: PennyLaneAI/pennylane
Pull request #5702 (2024-05-16 20:07:24)

Fix dimension ordering bug with sum_expand
**Context:**
`sum_expand` incorrectly order the dimensions when combining shot vectors, multiple measurements, and parameter broadcasting. Additionally, `qml.math.dot` raises an error when when only one of the operands is a scaler. For example:
```pycon
>>> import pennylane.numpy as np
>>> arr = np.array([1,2,3])
>>> qml.math.dot(arr, 2)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/astral.cai/Workspace/pennylane/pennylane/math/multi_dispatch.py", line 151, in wrapper
    return fn(*args, **kwargs)
  File "/Users/astral.cai/Workspace/pennylane/pennylane/math/multi_dispatch.py", line 365, in dot
    return np.tensordot(x, y, axes=[[-1], [-2]], like=like)
  File "/Users/astral.cai/Workspace/pennylane/venv/lib/python3.9/site-packages/autoray/autoray.py", line 80, in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
  File "/Users/astral.cai/Workspace/pennylane/pennylane/numpy/wrapper.py", line 117, in _wrapped
    res = obj(*args, **kwargs)
  File "/Users/astral.cai/Workspace/pennylane/venv/lib/python3.9/site-packages/autograd/tracer.py", line 48, in f_wrapped
    return f_raw(*args, **kwargs)
  File "/Users/astral.cai/Workspace/pennylane/venv/lib/python3.9/site-packages/numpy/core/numeric.py", line 1091, in tensordot
    if as_[axes_a[k]] != bs[axes_b[k]]:
IndexError: tuple index out of range
```

**Description of the Change:**
- Fixes the bug with `sum_expand`.
- Updates `dot` in to use `*` if either operand is a scaler.

**Related GitHub Issues:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/5700
[sc-63541]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5702?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.67%. Comparing base [(`5b81a8a`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/5b81a8a50a5ac6ee127d01a508ac014e654e2cbc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`7b311e3`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5702?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5702      +/-   ##
==========================================
- Coverage   99.68%   99.67%   -0.01%     
==========================================
  Files         415      415              
  Lines       38886    38590     -296     
==========================================
- Hits        38762    38465     -297     
- Misses        124      125       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5702?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index c1aeee482..c2525461f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -180,13 +180,21 @@
 * A correction is added to `bravyi_kitaev` to call the correct function for a FermiSentence input.
   [(#5671)](https://github.com/PennyLaneAI/pennylane/pull/5671)
 
+* Fixes a bug where `sum_expand` produces incorrect result dimensions when combining shot vectors, 
+  multiple measurements, and parameter broadcasting.
+  [(#5702)](https://github.com/PennyLaneAI/pennylane/pull/5702)
+
+* Fixes a bug in `qml.math.dot` that raises an error when only one of the operands is a scalar.
+  [(#5702)](https://github.com/PennyLaneAI/pennylane/pull/5702)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
 
 Lillian M. A. Frederiksen,
-Ahmed Darwish,
 Gabriel Bottrill,
+Astral Cai,
+Ahmed Darwish,
 Isaac De Vlugt,
 Pietropaolo Frisoni,
 Soran Jahangiri,
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index 47e44459d..6510bc995 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -317,8 +317,9 @@ def matmul(tensor1, tensor2, like=None):
 def dot(tensor1, tensor2, like=None):
     """Returns the matrix or dot product of two tensors.
 
-    * If both tensors are 0-dimensional, elementwise multiplication
-      is performed and a 0-dimensional scalar returned.
+    * If either tensor is 0-dimensional, elementwise multiplication
+      is performed and a 0-dimensional scalar or a tensor with the
+      same dimensions as the other tensor is returned.
 
     * If both tensors are 1-dimensional, the dot product is returned.
 
@@ -327,7 +328,7 @@ def dot(tensor1, tensor2, like=None):
 
     * If both tensors are 2-dimensional, the matrix product is returned.
 
-    * Finally, if the the first array is N-dimensional and the second array
+    * Finally, if the first array is N-dimensional and the second array
       M-dimensional, a sum product over the last dimension of the first array,
       and the second-to-last dimension of the second array is returned.
 
@@ -342,7 +343,7 @@ def dot(tensor1, tensor2, like=None):
 
     if like == "torch":
 
-        if x.ndim == 0 and y.ndim == 0:
+        if x.ndim == 0 or y.ndim == 0:
             return x * y
 
         if x.ndim <= 2 and y.ndim <= 2:
@@ -351,15 +352,17 @@ def dot(tensor1, tensor2, like=None):
         return np.tensordot(x, y, axes=[[-1], [-2]], like=like)
 
     if like in {"tensorflow", "autograd"}:
-        shape_y = len(np.shape(y))
-        shape_x = len(np.shape(x))
-        if shape_x == 0 and shape_y == 0:
+
+        ndim_y = len(np.shape(y))
+        ndim_x = len(np.shape(x))
+
+        if ndim_x == 0 or ndim_y == 0:
             return x * y
 
-        if shape_y == 1:
+        if ndim_y == 1:
             return np.tensordot(x, y, axes=[[-1], [0]], like=like)
 
-        if shape_x == 2 and shape_y == 2:
+        if ndim_x == 2 and ndim_y == 2:
             return x @ y
 
         return np.tensordot(x, y, axes=[[-1], [-2]], like=like)
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 8da3507e7..76b693f61 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -111,10 +111,7 @@ def naive_processing_fn(res, coeffs, offset):
     """
     dot_products = []
     for c, r in zip(coeffs, res):
-        if qml.math.ndim(c) == 0 and qml.math.size(r) != 1:
-            dot_products.append(qml.math.squeeze(r) * c)
-        else:
-            dot_products.append(qml.math.dot(qml.math.squeeze(r), c))
+        dot_products.append(qml.math.dot(qml.math.squeeze(r), c))
     if len(dot_products) == 0:
         return offset
     summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
@@ -344,8 +341,7 @@ def _sum_expand_processing_fn_grouping(
             coeffs.append(coeff)
         res_for_each_mp.append(naive_processing_fn(sub_res, coeffs, offset))
     if shots.has_partitioned_shots:
-        res_for_each_mp = qml.math.stack(res_for_each_mp, axis=0)
-        res_for_each_mp = qml.math.moveaxis(res_for_each_mp, 0, -1)
+        res_for_each_mp = qml.math.moveaxis(res_for_each_mp, 0, 1)
     return res_for_each_mp[0] if len(res_for_each_mp) == 1 else res_for_each_mp
 
 
@@ -368,9 +364,7 @@ def _sum_expand_processing_fn(
             coeffs.append(coeff)
         res_for_each_mp.append(naive_processing_fn(sub_res, coeffs, offset))
     if shots.has_partitioned_shots:
-        res_for_each_mp = qml.math.stack(res_for_each_mp, axis=0)
-        # Move dimensions around to make things work.
-        res_for_each_mp = qml.math.moveaxis(res_for_each_mp, 0, -1)
+        res_for_each_mp = qml.math.moveaxis(res_for_each_mp, 0, 1)
     return res_for_each_mp[0] if len(res_for_each_mp) == 1 else res_for_each_mp
 
 


Test patch:
diff --git a/tests/math/test_multi_dispatch.py b/tests/math/test_multi_dispatch.py
index 9f3d43b48..8eba39e0d 100644
--- a/tests/math/test_multi_dispatch.py
+++ b/tests/math/test_multi_dispatch.py
@@ -185,6 +185,7 @@ def test_gammainc(n, t, gamma_ref):
 
 
 def test_dot_autograd():
+
     x = np.array([1.0, 2.0], requires_grad=False)
     y = np.array([2.0, 3.0], requires_grad=True)
 
@@ -196,6 +197,50 @@ def test_dot_autograd():
     assert fn.allclose(qml_grad(fn.dot)(x, y), x)
 
 
+def test_dot_autograd_with_scalar():
+
+    x = np.array(1.0, requires_grad=False)
+    y = np.array([2.0, 3.0], requires_grad=True)
+
+    res = fn.dot(x, y)
+    assert isinstance(res, np.tensor)
+    assert res.requires_grad
+    assert fn.allclose(res, [2.0, 3.0])
+
+    res = fn.dot(y, x)
+    assert isinstance(res, np.tensor)
+    assert res.requires_grad
+    assert fn.allclose(res, [2.0, 3.0])
+
+
+def test_dot_tf_with_scalar():
+
+    x = tf.Variable(1.0)
+    y = tf.Variable([2.0, 3.0])
+
+    res = fn.dot(x, y)
+    assert isinstance(res, tf.Tensor)
+    assert fn.allclose(res, [2.0, 3.0])
+
+    res = fn.dot(y, x)
+    assert isinstance(res, tf.Tensor)
+    assert fn.allclose(res, [2.0, 3.0])
+
+
+def test_dot_torch_with_scalar():
+
+    x = torch.tensor(1.0)
+    y = torch.tensor([2.0, 3.0])
+
+    res = fn.dot(x, y)
+    assert isinstance(res, torch.Tensor)
+    assert fn.allclose(res, [2.0, 3.0])
+
+    res = fn.dot(y, x)
+    assert isinstance(res, torch.Tensor)
+    assert fn.allclose(res, [2.0, 3.0])
+
+
 def test_kron():
     """Test the kronecker product function."""
     x = torch.tensor([[1, 2], [3, 4]])
diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py
index 578430998..879c842e1 100644
--- a/tests/transforms/test_hamiltonian_expand.py
+++ b/tests/transforms/test_hamiltonian_expand.py
@@ -749,8 +749,7 @@ class TestSumExpand:
         res = circuit(theta)
 
         if isinstance(theta, np.ndarray):
-            expected = np.stack(expected).T
-            assert qml.math.shape(res) == (5, 4, 3)
+            assert qml.math.shape(res) == (5, 3, 4)
         else:
             assert qml.math.shape(res) == (5, 3)