Fix dimension ordering bug with sum_expand (#5702)

**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 #5700
[sc-63541]

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,

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)

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
 
 

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]])

tests/transforms/test_hamiltonian_expand.py

@@ -749,8 +749,7 @@ def circuit(x):
         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)