Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Composite operator parameters are not trainable
### Expected behavior

The parameters of composite ops are trainable (just like the 1st-generation Tensor and Hamiltonians).

### Actual behavior

They are not trainable.

### Additional information

This is because composite operators put their parameters in vanilla python lists, and those are considered not trainable. My code snippet first shows a full example (mimicking [how pennylane currently sets trainable parameters](https://github.com/PennyLaneAI/pennylane/blob/601c5ca6f5caec9c917c2c57fb5b61f6ed5f4ffa/pennylane/interfaces/autograd.py#L72-L73)), and then a minimal demonstration of what we currently consider trainable.

### Source code

```shell
dev = qml.device("default.qubit", wires=2)
@qml.qnode(dev)
def circuit():
    return qml.expval(qml.prod(qml.PauliX(0), qml.RX(qnp.array(1.1), 1)))

circuit()
params = circuit.tape.get_parameters(trainable_only=False)
print(params, qml.math.get_trainable_indices(params))
# [[], [tensor(1.1, requires_grad=True)]] set()


### Minimal demo of why it doesn't work
from pennylane.math import requires_grad; from pennylane import numpy as np
print(requires_grad(np.array(1)), requires_grad([np.array(1)]))
# True False
```


### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.30.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/matthews/.pyenv/versions/3.9.13/envs/pl/lib/python3.9/site-packages
Editable project location: /Users/matthews/src/github.com/PennyLaneAI/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning

Platform info:           macOS-13.2.1-arm64-arm-64bit
Python version:          3.9.13
Numpy version:           1.23.5
Scipy version:           1.10.0
Installed devices:
- default.gaussian (PennyLane-0.30.0.dev0)
- default.mixed (PennyLane-0.30.0.dev0)
- default.qubit (PennyLane-0.30.0.dev0)
- default.qubit.autograd (PennyLane-0.30.0.dev0)
- default.qubit.jax (PennyLane-0.30.0.dev0)
- default.qubit.tf (PennyLane-0.30.0.dev0)
- default.qubit.torch (PennyLane-0.30.0.dev0)
- default.qutrit (PennyLane-0.30.0.dev0)
- null.qubit (PennyLane-0.30.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.29.0.dev11)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request:
Update ScalarSymbolicOps and trainability
**Context:**
We want all operator data to be as flat as possible. Also, I had previously done this for Sum and Prod, but didn't test the consequential result which is that their data can now be considered trainable.

**Description of the Change:**
1. Move identical data logic from SProd and Exp to the parent class, ScalarSymbolicOp. Note that while doing this, I also made the data for `Exp` flat.
2. Add a custom `data` property definition to Pow so it keeps its old behaviour of treating `z` as a hyperparameter (eg. excluded from `op.data`)
3. The `copy` dunder for Exp was moved to the parent class so its siblings can have it too : )
4. Add tests to ensure Exp, Sum and Prod data is trainable (there was already an SProd test that actually measures the gradient)

**Benefits:**
- It is easier to detect trainability of flat data, and more consistent with the rest of PennyLane
- Less copied code
- The new tape converter integration can move forward thanks to the addition of a data setter to all ScalarSymbolicOps

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
#3958, #3899, fixes #3916

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 32a1ddcfa..74af67582 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -158,8 +158,9 @@
 * `retworkx` has been renamed to `rustworkx` to accomodate the change in name for the package.
   [(#3975)](https://github.com/PennyLaneAI/pennylane/pull/3975)
 
-* `Sum`, `Prod`, and `SProd` operator data is now a flat list, instead of nested.
+* `Exp`, `Sum`, `Prod`, and `SProd` operator data is now a flat list, instead of nested.
   [(#3958)](https://github.com/PennyLaneAI/pennylane/pull/3958)
+  [(#3983)](https://github.com/PennyLaneAI/pennylane/pull/3983)
 
 * `qml.transforms.convert_to_numpy_parameters` is added to convert a circuit with interface-specific parameters to one
   with only numpy parameters. This transform is designed to replace `qml.tape.Unwrap`.
@@ -195,8 +196,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.
+* `Exp`, `Sum`, `Prod`, and `SProd` operator data is now a flat list, instead of nested.
   [(#3958)](https://github.com/PennyLaneAI/pennylane/pull/3958)
+  [(#3983)](https://github.com/PennyLaneAI/pennylane/pull/3983)
 
 <h3>Deprecations</h3>
 
diff --git a/pennylane/ops/op_math/evolution.py b/pennylane/ops/op_math/evolution.py
index 9092e5853..84d021e2d 100644
--- a/pennylane/ops/op_math/evolution.py
+++ b/pennylane/ops/op_math/evolution.py
@@ -15,6 +15,7 @@
 This submodule defines the Evolution class.
 """
 import warnings
+from copy import copy
 from warnings import warn
 
 import pennylane as qml
@@ -93,6 +94,14 @@ class Evolution(Exp):
             else f"Evolution({self.coeff} {self.base.name})"
         )
 
+    @property
+    def data(self):
+        return self._data
+
+    @data.setter
+    def data(self, new_data):
+        self._data = new_data
+
     @property
     def param(self):
         """A real coefficient with ``1j`` factored out."""
@@ -145,3 +154,8 @@ class Evolution(Exp):
                 f"The generator is not defined."
             )
         return self.base
+
+    def __copy__(self):
+        copied = super().__copy__()
+        copied._data = copy(self._data)
+        return copied
diff --git a/pennylane/ops/op_math/exp.py b/pennylane/ops/op_math/exp.py
index ce2de0a78..0ae1fc640 100644
--- a/pennylane/ops/op_math/exp.py
+++ b/pennylane/ops/op_math/exp.py
@@ -14,7 +14,6 @@
 """
 This submodule defines the symbolic operation that stands for an exponential of an operator.
 """
-from copy import copy
 from typing import List
 from warnings import warn
 
@@ -173,7 +172,6 @@ class Exp(ScalarSymbolicOp, Operation):
     # pylint: disable=too-many-arguments
     def __init__(self, base, coeff=1, num_steps=None, do_queue=True, id=None):
         super().__init__(base, scalar=coeff, do_queue=do_queue, id=id)
-        self._data = [[coeff], self.base.data]
         self.grad_recipe = [None]
         self.num_steps = num_steps
 
@@ -184,31 +182,10 @@ class Exp(ScalarSymbolicOp, Operation):
             else f"Exp({self.coeff} {self.base.name})"
         )
 
-    # pylint: disable=attribute-defined-outside-init
-    def __copy__(self):
-        # this method needs to be overwritten because the base must be copied too.
-        copied_op = object.__new__(type(self))
-        # copied_op must maintain inheritance structure of self
-        # Relevant for symbolic ops that mix in operation-specific components.
-
-        for attr, value in vars(self).items():
-            if attr not in {"_hyperparameters"}:
-                setattr(copied_op, attr, value)
-
-        copied_op._hyperparameters = copy(self.hyperparameters)
-        copied_op.hyperparameters["base"] = copy(self.base)
-        copied_op._data = copy(self._data)
-
-        return copied_op
-
     @property
     def hash(self):
         return hash((str(self.name), self.base.hash, str(self.coeff)))
 
-    @property
-    def data(self):
-        return self._data
-
     @property
     def coeff(self):
         """The numerical coefficient of the operator in the exponent."""
diff --git a/pennylane/ops/op_math/pow.py b/pennylane/ops/op_math/pow.py
index 4d32b4673..b2a78776f 100644
--- a/pennylane/ops/op_math/pow.py
+++ b/pennylane/ops/op_math/pow.py
@@ -236,6 +236,15 @@ class Pow(ScalarSymbolicOp):
     def ndim_params(self):
         return self.base.ndim_params
 
+    @property
+    def data(self):
+        """The trainable parameters"""
+        return self.base.data
+
+    @data.setter
+    def data(self, new_data):
+        self.base.data = new_data
+
     def label(self, decimals=None, base_label=None, cache=None):
         z_string = format(self.z).translate(_superscript)
         base_label = self.base.label(decimals, base_label, cache=cache)
diff --git a/pennylane/ops/op_math/sprod.py b/pennylane/ops/op_math/sprod.py
index 6e03cd4c2..839dd5466 100644
--- a/pennylane/ops/op_math/sprod.py
+++ b/pennylane/ops/op_math/sprod.py
@@ -163,17 +163,6 @@ class SProd(ScalarSymbolicOp):
 
         return base_label or f"{scalar_val}*{self.base.label(decimals=decimals, cache=cache)}"
 
-    @property
-    def data(self):
-        """The trainable parameters"""
-        return [self.scalar, *self.base.data]
-
-    @data.setter
-    def data(self, new_data):
-        self.scalar = new_data[0]
-        if len(new_data) > 1:
-            self.base.data = new_data[1:]
-
     @property
     def num_params(self):
         """Number of trainable parameters that the operator depends on.
diff --git a/pennylane/ops/op_math/symbolicop.py b/pennylane/ops/op_math/symbolicop.py
index a109e5a3b..235f2d0b1 100644
--- a/pennylane/ops/op_math/symbolicop.py
+++ b/pennylane/ops/op_math/symbolicop.py
@@ -173,6 +173,15 @@ class ScalarSymbolicOp(SymbolicOp):
     def batch_size(self):
         return self._batch_size
 
+    @property
+    def data(self):
+        return [self.scalar, *self.base.data]
+
+    @data.setter
+    def data(self, new_data):
+        self.scalar = new_data[0]
+        self.base.data = new_data[1:]
+
     def _check_and_compute_batch_size(self, scalar):
         batched_scalar = qml.math.ndim(scalar) > 0
         scalar_size = qml.math.size(scalar)


Test patch:
diff --git a/tests/ops/op_math/test_evolution.py b/tests/ops/op_math/test_evolution.py
index e814b1dfb..42a754ab2 100644
--- a/tests/ops/op_math/test_evolution.py
+++ b/tests/ops/op_math/test_evolution.py
@@ -83,6 +83,13 @@ class TestEvolution:
         assert op.coeff == 1j * op.data[0]
         assert op.param == op.data[0]
 
+        new_param = np.array(2.345)
+        op.data = [new_param]
+
+        assert op.data == [new_param]
+        assert op.coeff == 1j * op.data[0]
+        assert op.data == op.data[0]
+
     def test_repr_paulix(self):
         """Test the __repr__ method when the base is a simple observable."""
         op = Evolution(qml.PauliX(0), 3)
diff --git a/tests/ops/op_math/test_exp.py b/tests/ops/op_math/test_exp.py
index c28513970..616931229 100644
--- a/tests/ops/op_math/test_exp.py
+++ b/tests/ops/op_math/test_exp.py
@@ -44,8 +44,8 @@ class TestInitialization:
         assert op.id == "something"
 
         assert op.num_params == 1
-        assert op.parameters == [[1], []]
-        assert op.data == [[1], []]
+        assert op.parameters == [1]
+        assert op.data == [1]
 
         assert op.wires == qml.wires.Wires("a")
 
@@ -63,8 +63,8 @@ class TestInitialization:
         assert op.name == "Exp"
 
         assert op.num_params == 1
-        assert op.parameters == [[coeff], []]
-        assert op.data == [[coeff], []]
+        assert op.parameters == [coeff]
+        assert op.data == [coeff]
 
         assert op.wires == qml.wires.Wires(("b", "c"))
 
@@ -82,7 +82,7 @@ class TestInitialization:
         assert op.name == "Exp"
 
         assert op.num_params == 2
-        assert op.data == [coeff, [base_coeff]]
+        assert op.data == [coeff, base_coeff]
 
         assert op.wires == qml.wires.Wires(5)
 
@@ -111,7 +111,15 @@ class TestProperties:
         base = qml.RX(phi, wires=0)
         op = Exp(base, coeff)
 
-        assert op.data == [[coeff], [phi]]
+        assert op.data == [coeff, phi]
+
+        new_phi = np.array(0.1234)
+        new_coeff = np.array(3.456)
+        op.data = [new_coeff, new_phi]
+
+        assert op.data == [new_coeff, new_phi]
+        assert op.base.data == [new_phi]
+        assert op.scalar == new_coeff
 
     # pylint: disable=protected-access
     def test_queue_category_ops(self):
@@ -933,3 +941,16 @@ class TestDifferentiation:
             op1.parameter_frequencies()
 
         assert op2.parameter_frequencies == [(4.0,)]
+
+    def test_params_can_be_considered_trainable(self):
+        """Tests that the parameters of an Exp are considered trainable."""
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(x, coeff):
+            Exp(qml.RX(x, 0), coeff)
+            return qml.expval(qml.PauliZ(0))
+
+        with pytest.warns(UserWarning):
+            circuit(np.array(2.0), np.array(0.5))
+        assert circuit.tape.trainable_params == [0, 1]
diff --git a/tests/ops/op_math/test_prod.py b/tests/ops/op_math/test_prod.py
index c2dea2759..829cd35a0 100644
--- a/tests/ops/op_math/test_prod.py
+++ b/tests/ops/op_math/test_prod.py
@@ -1300,6 +1300,18 @@ class TestIntegration:
         res2 = batched_no_prod(x, y)
         assert qml.math.allclose(res1, res2)
 
+    def test_params_can_be_considered_trainable(self):
+        """Tests that the parameters of a Prod are considered trainable."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.expval(qml.prod(qml.RX(1.1, 0), qml.RY(qnp.array(2.2), 1)))
+
+        with pytest.warns(UserWarning):
+            circuit()
+        assert circuit.tape.trainable_params == [1]
+
 
 class TestSortWires:
     """Tests for the wire sorting algorithm."""
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index cc094c90b..ed633bb0f 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -1076,6 +1076,18 @@ class TestIntegration:
         with pytest.warns(UserWarning, match="Sum might not be hermitian."):
             my_circ()
 
+    def test_params_can_be_considered_trainable(self):
+        """Tests that the parameters of a Sum are considered trainable."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.expval(Sum(qml.RX(1.1, 0), qml.RY(qnp.array(2.2), 0)))
+
+        with pytest.warns(UserWarning):
+            circuit()
+        assert circuit.tape.trainable_params == [1]
+
 
 # pylint: disable=too-few-public-methods
 class TestArithmetic:
diff --git a/tests/ops/op_math/test_symbolic_op.py b/tests/ops/op_math/test_symbolic_op.py
index d5dd9d8f4..74058ec4b 100644
--- a/tests/ops/op_math/test_symbolic_op.py
+++ b/tests/ops/op_math/test_symbolic_op.py
@@ -18,7 +18,7 @@ import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.ops.op_math import SymbolicOp
+from pennylane.ops.op_math import SymbolicOp, ScalarSymbolicOp
 from pennylane.wires import Wires
 
 
@@ -26,6 +26,11 @@ class TempOperator(qml.operation.Operator):
     num_wires = 1
 
 
+class TempScalar(ScalarSymbolicOp):
+    def _matrix(scalar, mat):
+        pass
+
+
 def test_intialization():
     """Test initialization for a SymbolicOp"""
     base = TempOperator("a")
@@ -203,3 +208,41 @@ class TestQueuing:
             op = SymbolicOp(base, do_queue=False)
 
         assert len(q) == 0
+
+
+class TestScalarSymbolicOp:
+    """Tests for the ScalarSymbolicOp class."""
+
+    def test_init(self):
+        base = TempOperator(1.1, wires=[0])
+        scalar = 2.2
+        op = TempScalar(base, scalar)
+        assert isinstance(op.scalar, float)
+        assert op.scalar == 2.2
+        assert op.data == [2.2, 1.1]
+
+        base = TempOperator(1.1, wires=[0])
+        scalar = [2.2, 3.3]
+        op = TempScalar(base, scalar)
+        assert isinstance(op.scalar, np.ndarray)
+        assert np.all(op.scalar == [2.2, 3.3])
+        assert np.all(op.data[0] == op.scalar)
+        assert op.data[1] == 1.1
+
+    def test_data(self):
+        """Tests the data property."""
+        op = TempScalar(TempOperator(1.1, wires=[0]), 2.2)
+        assert op.scalar == 2.2
+        assert op.data == [2.2, 1.1]
+
+        # check setting through ScalarSymbolicOp
+        op.data = [3.3, 4.4]
+        assert op.data == [3.3, 4.4]
+        assert op.scalar == 3.3
+        assert op.base.data == [4.4]
+
+        # check setting through base
+        op.base.data = [5.5]
+        assert op.data == [3.3, 5.5]
+        assert op.scalar == 3.3
+        assert op.base.data == [5.5]