Repository: PennyLaneAI/pennylane
Pull request #3958 (2023-04-03 14:31:02)

flatten Composite and SProd data
**Context:**
Composite and SProd data is not flat, and this is inconsistent with the rest of PennyLane operators. It's also hard to work with unless you know the expected form.

**Description of the Change:**
Flatten data for those operators.
1. Composite op data will just be a flat list of all data for all operands
2. SProd data will just be `[self.scalar] + self.base.data`

**Benefits:**
- Composite and SProd operator data is easier to work with
- It is slightly more performant when setting and getting data

**Possible Drawbacks:**
- Not a drawback for this PR, but a note. Composite op data getting and setting is still not performant enough. It is better to avoid doing that altogether
- Some people might like the old nested way? I doubt this - especially since Sums will most often represent Hamiltonians, and that data will usually look like a flat list of coeffs.

**Related GitHub Issues:**
Would help unblock #3899 in a clean way

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://codecov.io/gh/PennyLaneAI/pennylane/pull/3958?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3958](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3958?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (dd291db) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/270a453dc7162afa05cc34e4d577adc029ec9eb9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (270a453) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3958   +/-   ##
=======================================
  Coverage   99.73%   99.73%           
=======================================
  Files         345      345           
  Lines       30197    30202    +5     
=======================================
+ Hits        30118    30123    +5     
  Misses         79       79           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3958?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/ops/op\_math/composite.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3958?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2NvbXBvc2l0ZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/op\_math/sprod.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3958?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3Nwcm9kLnB5) | `96.70% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/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 77184f803..ef4cb5940 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -141,6 +141,9 @@
   This allows devices that inherit from `QubitDevice` to override and customize their definition of diagonalizing gates.
   [(#3938)](https://github.com/PennyLaneAI/pennylane/pull/3938)
 
+* `Sum`, `Prod`, and `SProd` operator data is now a flat list, instead of nested.
+  [(#3958)](https://github.com/PennyLaneAI/pennylane/pull/3958)
+
 <h3>Breaking changes</h3>
 
 * Both JIT interfaces are not compatible with Jax `>0.4.3`, we raise an error for those versions.
@@ -167,6 +170,9 @@
   For example, you can no longer create `StateMP(qml.PauliX(0))` or `PurityMP(eigvals=(-1,1), wires=Wires(0))`.
   [(#3898)](https://github.com/PennyLaneAI/pennylane/pull/3898)
 
+* `Sum`, `Prod`, and `SProd` operator data is now a flat list, instead of nested.
+  [(#3958)](https://github.com/PennyLaneAI/pennylane/pull/3958)
+
 <h3>Deprecations</h3>
 
 <h3>Documentation</h3>
diff --git a/pennylane/ops/op_math/composite.py b/pennylane/ops/op_math/composite.py
index 2653110b5..a5361c586 100644
--- a/pennylane/ops/op_math/composite.py
+++ b/pennylane/ops/op_math/composite.py
@@ -112,13 +112,16 @@ class CompositeOp(Operator):
     @property
     def data(self):
         """Create data property"""
-        return [op.data for op in self]
+        return [d for op in self for d in op.data]
 
     @data.setter
     def data(self, new_data):
         """Set the data property"""
-        for new_entry, op in zip(new_data, self):
-            op.data = new_entry
+        for op in self:
+            op_num_params = op.num_params
+            if op_num_params > 0:
+                op.data = new_data[:op_num_params]
+                new_data = new_data[op_num_params:]
 
     @property
     def num_wires(self):
diff --git a/pennylane/ops/op_math/sprod.py b/pennylane/ops/op_math/sprod.py
index 5b21f30fa..6e03cd4c2 100644
--- a/pennylane/ops/op_math/sprod.py
+++ b/pennylane/ops/op_math/sprod.py
@@ -166,13 +166,13 @@ class SProd(ScalarSymbolicOp):
     @property
     def data(self):
         """The trainable parameters"""
-        return [[self.scalar], self.base.data]  # Not sure if this is the best way to deal with this
+        return [self.scalar, *self.base.data]
 
     @data.setter
     def data(self, new_data):
-        self.scalar = new_data[0][0]
+        self.scalar = new_data[0]
         if len(new_data) > 1:
-            self.base.data = new_data[1]
+            self.base.data = new_data[1:]
 
     @property
     def num_params(self):


Test patch:
diff --git a/tests/ops/op_math/test_composite.py b/tests/ops/op_math/test_composite.py
index 8d4e6dacd..49579aaa0 100644
--- a/tests/ops/op_math/test_composite.py
+++ b/tests/ops/op_math/test_composite.py
@@ -93,25 +93,26 @@ class TestConstruction:
 
     def test_parameters(self):
         """Test that parameters are initialized correctly."""
-        op = ValidOp(qml.RX(9.87, wires=0), qml.Rot(1.23, 4.0, 5.67, wires=1))
-        assert op.parameters == [[9.87], [1.23, 4.0, 5.67]]
+        op = ValidOp(qml.RX(9.87, wires=0), qml.Rot(1.23, 4.0, 5.67, wires=1), qml.PauliX(0))
+        assert op.parameters == [9.87, 1.23, 4.0, 5.67]
 
     def test_data(self):
         """Test that data is initialized correctly."""
-        op = ValidOp(qml.RX(9.87, wires=0), qml.Rot(1.23, 4.0, 5.67, wires=1))
-        assert op.data == [[9.87], [1.23, 4.0, 5.67]]
+        op = ValidOp(qml.RX(9.87, wires=0), qml.Rot(1.23, 4.0, 5.67, wires=1), qml.PauliX(0))
+        assert op.data == [9.87, 1.23, 4.0, 5.67]
 
     def test_data_setter(self):
         """Test the setter method for data"""
-        op = ValidOp(qml.RX(9.87, wires=0), qml.Rot(1.23, 4.0, 5.67, wires=1))
-        assert op.data == [[9.87], [1.23, 4.0, 5.67]]
+        op = ValidOp(qml.RX(9.87, wires=0), qml.Rot(1.23, 4.0, 5.67, wires=1), qml.PauliX(0))
+        assert op.data == [9.87, 1.23, 4.0, 5.67]
 
-        new_data = [[1.23], [0.0, -1.0, -2.0]]
+        new_data = [1.23, 0.0, -1.0, -2.0]
         op.data = new_data
         assert op.data == new_data
 
-        for op, new_entry in zip(op.operands, new_data):
-            assert op.data == new_entry
+        for o in op:
+            assert o.data == new_data[: o.num_params]
+            new_data = new_data[o.num_params :]
 
     def test_ndim_params_raises_error(self):
         """Test that calling ndim_params raises a ValueError."""
diff --git a/tests/ops/op_math/test_prod.py b/tests/ops/op_math/test_prod.py
index 344d935b2..c2dea2759 100644
--- a/tests/ops/op_math/test_prod.py
+++ b/tests/ops/op_math/test_prod.py
@@ -129,8 +129,8 @@ class TestInitialization:
         assert prod_op.id == id
         assert prod_op.queue_idx is None
 
-        assert prod_op.data == [[], [0.23]]
-        assert prod_op.parameters == [[], [0.23]]
+        assert prod_op.data == [0.23]
+        assert prod_op.parameters == [0.23]
         assert prod_op.num_params == 1
 
     def test_hash(self):
diff --git a/tests/ops/op_math/test_sprod.py b/tests/ops/op_math/test_sprod.py
index aefb03734..e9e7beb62 100644
--- a/tests/ops/op_math/test_sprod.py
+++ b/tests/ops/op_math/test_sprod.py
@@ -112,31 +112,49 @@ class TestInitialization:
         assert sprod_op.id == test_id
         assert sprod_op.queue_idx is None
 
-        assert sprod_op.data == [[3.14], [0.23]]
-        assert sprod_op.parameters == [[3.14], [0.23]]
+        assert sprod_op.data == [3.14, 0.23]
+        assert sprod_op.parameters == [3.14, 0.23]
         assert sprod_op.num_params == 2
 
     def test_parameters(self):
         sprod_op = s_prod(9.87, qml.Rot(1.23, 4.0, 5.67, wires=1))
-        assert sprod_op.parameters == [[9.87], [1.23, 4.0, 5.67]]
+        assert sprod_op.parameters == [9.87, 1.23, 4.0, 5.67]
 
     def test_data(self):
         sprod_op = s_prod(9.87, qml.Rot(1.23, 4.0, 5.67, wires=1))
-        assert sprod_op.data == [[9.87], [1.23, 4.0, 5.67]]
+        assert sprod_op.data == [9.87, 1.23, 4.0, 5.67]
 
     def test_data_setter(self):
         """Test the setter method for data"""
         scalar, angles = (9.87, (1.23, 4.0, 5.67))
-        old_data = [[9.87], [1.23, 4.0, 5.67]]
+        old_data = [9.87, 1.23, 4.0, 5.67]
 
         sprod_op = s_prod(scalar, qml.Rot(*angles, wires=1))
         assert sprod_op.data == old_data
 
-        new_data = [[1.23], [0.0, -1.0, -2.0]]
+        new_data = [1.23, 0.0, -1.0, -2.0]
         sprod_op.data = new_data
         assert sprod_op.data == new_data
-        assert sprod_op.scalar == new_data[0][0]
-        assert sprod_op.base.data == new_data[1]
+        assert sprod_op.scalar == new_data[0]
+        assert sprod_op.base.data == new_data[1:]
+
+    def test_data_setter_shallow(self):
+        """Test the setter method for data with a non-parametric base op."""
+        op = s_prod(0.1, qml.PauliX(0))
+        op.data = [0.2]
+        assert op.data == [0.2] == [op.scalar]
+
+    def test_data_setter_deep(self):
+        """Test the setter method for data with a deep base operator."""
+        op = s_prod(0.1, qml.sum(qml.PauliX(0), qml.prod(qml.PauliY(0), qml.RX(0.2, 1))))
+        assert op.data == [0.1, 0.2]
+
+        new_data = [0.3, 0.4]
+        op.data = new_data
+        assert op.data == new_data
+        assert op.scalar == 0.3
+        assert op.base[1].data == [0.4]
+        assert op.base[1][1].data == [0.4]
 
     @pytest.mark.parametrize("scalar, op", ops)
     def test_terms(self, op, scalar):
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index 1aa2fa53a..cc094c90b 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -126,8 +126,8 @@ class TestInitialization:
         if sum_method.__name__ == sum.__name__:
             assert sum_op.id == id
 
-        assert sum_op.data == [[], [0.23]]
-        assert sum_op.parameters == [[], [0.23]]
+        assert sum_op.data == [0.23]
+        assert sum_op.parameters == [0.23]
         assert sum_op.num_params == 1
 
     @pytest.mark.parametrize("sum_method", [sum_using_dunder_method, qml.sum])
@@ -142,8 +142,8 @@ class TestInitialization:
         assert sum_op.name == "Sum"
         assert sum_op.id is None
 
-        assert sum_op.data == [[[], [0.23]], [9.87]]
-        assert sum_op.parameters == [[[], [0.23]], [9.87]]
+        assert sum_op.data == [0.23, 9.87]
+        assert sum_op.parameters == [0.23, 9.87]
         assert sum_op.num_params == 2
 
     @pytest.mark.parametrize("ops_lst", ops)
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
index c52124c45..c6d5554b1 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -189,7 +189,7 @@ class TestQubitUnitaryXYXDecomposition:
             assert obtained_gates[i].wires == Wires("a"), "Incorrect wire"
         # Check the global phase
         assert qml.math.isclose(
-            qml.math.unwrap(obtained_gates[3].parameters[0]), expected_params[3]
+            qml.math.unwrap(obtained_gates[3].parameters[:1]), expected_params[3]
         ), "Incorrect global phase"
         # Now we check the XYX rotation angles
         assert qml.math.allclose(