Repository: PennyLaneAI/pennylane
Pull request #3899 (2023-03-14 15:01:04)

Add function for convert circuit parameters to numpy
Currently, we use the `qml.tape.Unwrap` context manager to convert parameters to numpy before sending them to a device.

This context manager modifies the input in-place and resets it at the end. It uses `QuantumScript.set_parameters`, which iterates over all operations, modifies them in place, and then checks batching.

Notably, it calls `op._check_batching` for every single parameter that gets set, not every operator that has parameter set.

This causes massive performance problems for large observables. We are also setting the parameters for operators that might not even need it.


This PR proposes a transform `convert_to_numpy_parameters` that creates a new quantum script with numpy or built-ins parameters.  Currently, it is in the `devices.preprocessing` module, as it was designed thinking about new devices, and making it easy for new devices to convert to vanilla numpy parameters.

Given the huge performance problems with `Unwrap`, we should consider using this everywhere instead.


![Screenshot 2023-03-14 at 10 48 22 AM](https://user-images.githubusercontent.com/6364575/225042785-b5a0541d-4979-4609-b012-9be084b25241.png)

When using a `Sum` object instead of a Hamiltonian, `Unwrap` takes 1min 36seconds, a great deal more time.  `convert_to_numpy_parameters` takes 8.42 ms, twice the amount of time as for `Hamiltonian`, but still substantially faster than the `Unwrap` case.  

If we are going to immediately replace `Unwrap` with this instead, we should better consider the location for this conversion utility.



codecov[bot]: ## [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3899?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3899](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3899?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (898452d) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/7c8859acfad4cede87859d1f8d25448c97dc3147?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7c8859a) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3899   +/-   ##
=======================================
  Coverage   99.70%   99.70%           
=======================================
  Files         342      344    +2     
  Lines       29853    29877   +24     
=======================================
+ Hits        29766    29790   +24     
  Misses         87       87           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3899?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/preprocessing/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3899?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcHJlcHJvY2Vzc2luZy9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [...ne/devices/preprocessing/internal\_preprocessing.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3899?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcHJlcHJvY2Vzc2luZy9pbnRlcm5hbF9wcmVwcm9jZXNzaW5nLnB5) | `100.00% <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)

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.
timmysilv: it's as if you magically looked at my local branch and nabbed the changes I made! looks excellent. I want to confirm a thing or two before approving, but I will likely approve today.
albi3ro: > Thanks Christina! I think that this looks good. Is there a timeline for when this would replace `tape.unwrap`?

@mudit2812  We are working on a PR to replace the use of `unwrap` in the interfaces module. 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fafd8d3d4..7b37cb752 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -156,6 +156,10 @@
 * `Sum`, `Prod`, and `SProd` operator data is now a flat list, instead of nested.
   [(#3958)](https://github.com/PennyLaneAI/pennylane/pull/3958)
 
+* `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`.
+  [(#3899)](https://github.com/PennyLaneAI/pennylane/pull/3899)
+
 * `qml.operation.WiresEnum.AllWires` is now -2 instead of 0 to avoid the
   ambiguity between `op.num_wires = 0` and `op.num_wires = AllWires`.
   [(#3978)](https://github.com/PennyLaneAI/pennylane/pull/3978)
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 2825a60e1..156076fed 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -152,6 +152,13 @@ more tapes as well as a classical processing function.
     ~transforms.sign_expand
     ~transforms.sum_expand
 
+This transform accepts a single tape and returns a single tape:
+
+.. autosummary::
+    :toctree: api
+
+    ~transforms.convert_to_numpy_parameters
+
 Decorators and utility functions
 --------------------------------
 
@@ -185,6 +192,7 @@ Transforms for error mitigation
     ~transforms.fold_global
     ~transforms.poly_extrapolate
     ~transforms.richardson_extrapolate
+
 """
 # Import the decorators first to prevent circular imports when used in other transforms
 from .batch_transform import batch_transform, map_batch_transform
@@ -195,6 +203,7 @@ from .batch_input import batch_input
 from .batch_partial import batch_partial
 from .classical_jacobian import classical_jacobian
 from .condition import cond, Conditional
+from .convert_to_numpy_parameters import convert_to_numpy_parameters
 from .compile import compile
 from .decompositions import zyz_decomposition, xyx_decomposition, two_qubit_decomposition
 from .defer_measurements import defer_measurements
diff --git a/pennylane/transforms/convert_to_numpy_parameters.py b/pennylane/transforms/convert_to_numpy_parameters.py
new file mode 100644
index 000000000..3b7dad410
--- /dev/null
+++ b/pennylane/transforms/convert_to_numpy_parameters.py
@@ -0,0 +1,93 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This file contains preprocessings steps that may be called internally
+during execution.
+"""
+import copy
+
+import pennylane as qml
+from pennylane import math
+from pennylane.tape import QuantumScript
+
+
+def _convert_op_to_numpy_data(op: qml.operation.Operator) -> qml.operation.Operator:
+    if math.get_interface(*op.data) == "numpy":
+        return op
+    # Use operator method to change parameters when it become available
+    copied_op = copy.copy(op)
+    copied_op.data = math.unwrap(op.data)
+    return copied_op
+
+
+def _convert_measurement_to_numpy_data(
+    m: qml.measurements.MeasurementProcess,
+) -> qml.measurements.MeasurementProcess:
+    if m.obs is None or math.get_interface(*m.obs.data) == "numpy":
+        return m
+    # Use measurement method to change parameters when it becomes available
+    copied_m = copy.copy(m)
+    if isinstance(copied_m.obs, qml.operation.Tensor):
+        copied_m.obs.data = math.unwrap([o.data for o in m.obs.obs])
+    else:
+        copied_m.obs.data = math.unwrap(m.obs.data)
+    return copied_m
+
+
+# pylint: disable=protected-access
+def convert_to_numpy_parameters(circuit: QuantumScript) -> QuantumScript:
+    """Transforms a circuit to one with purely numpy parameters.
+
+    Args:
+        circuit (QuantumScript): a circuit with parameters of any interface
+
+    Returns:
+        QuantumScript: A circuit with purely numpy parameters
+
+    .. seealso::
+
+        :class:`pennylane.tape.Unwrap` modifies a :class:`~.pennylane.tape.QuantumScript` in place instead of creating
+        a new class. It will also set all parameters on the circuit, not just ones that need to be unwrapped.
+
+    >>> ops = [qml.S(0), qml.RX(torch.tensor(0.1234), 0)]
+    >>> measurements = [qml.state(), qml.expval(qml.Hermitian(torch.eye(2), 0))]
+    >>> circuit = qml.tape.QuantumScript(ops, measurements )
+    >>> new_circuit = convert_to_numpy_parameters(circuit)
+    >>> new_circuit.circuit
+    [S(wires=[0]),
+    RX(0.1234000027179718, wires=[0]),
+    state(wires=[]),
+    expval(Hermitian(array([[1., 0.],
+            [0., 1.]], dtype=float32), wires=[0]))]
+
+    If the component's data does not need to be transformed, it is left uncopied.
+
+    >>> circuit[0] is new_circuit[0]
+    True
+    >>> circuit[1] is new_circuit[1]
+    False
+    >>> circuit[2] is new_circuit[2]
+    True
+    >>> circuit[3] is new_circuit[3]
+    False
+
+    """
+    new_prep = (_convert_op_to_numpy_data(op) for op in circuit._prep)
+    new_ops = (_convert_op_to_numpy_data(op) for op in circuit._ops)
+    new_measurements = (_convert_measurement_to_numpy_data(m) for m in circuit.measurements)
+    new_circuit = circuit.__class__(new_ops, new_measurements, new_prep)
+    # must preserve trainable params as we lose information about the machine learning interface
+    new_circuit.trainable_params = circuit.trainable_params
+    new_circuit._qfunc_output = circuit._qfunc_output
+    return new_circuit


Test patch:
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index 329af8308..e1c400221 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -24,6 +24,7 @@ tests/transforms/test_sign_expand.py
 tests/ops/qubit/test_special_unitary.py
 tests/ops/qubit/test_matrix_ops.py
 tests/ops/op_math/test_exp.py
+tests/transforms/test_convert_to_numpy_parameters.py
 tests/transforms/test_decompositions.py
 tests/transforms/test_qfunc_transform.py
 tests/ops/op_math/test_exp.py
diff --git a/tests/transforms/test_convert_to_numpy_parameters.py b/tests/transforms/test_convert_to_numpy_parameters.py
new file mode 100644
index 000000000..b61ebb91d
--- /dev/null
+++ b/tests/transforms/test_convert_to_numpy_parameters.py
@@ -0,0 +1,109 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""This file tests the convert_to_numpy_parameters function."""
+import pytest
+
+import numpy as np
+
+import pennylane as qml
+from pennylane.transforms import convert_to_numpy_parameters
+from pennylane.transforms.convert_to_numpy_parameters import (
+    _convert_op_to_numpy_data,
+    _convert_measurement_to_numpy_data,
+)
+
+
+ml_frameworks_list = [
+    pytest.param("autograd", marks=pytest.mark.autograd),
+    pytest.param("jax", marks=pytest.mark.jax),
+    pytest.param("torch", marks=pytest.mark.torch),
+    pytest.param("tensorflow", marks=pytest.mark.tf),
+]
+
+
+@pytest.mark.parametrize("framework", ml_frameworks_list)
+def test_convert_arrays_to_numpy(framework):
+    """Tests that convert_to_numpy_parameters works with autograd arrays."""
+
+    x = qml.math.asarray(np.array(1.234), like=framework)
+    y = qml.math.asarray(np.array(0.652), like=framework)
+    M = qml.math.asarray(np.eye(2), like=framework)
+    state = qml.math.asarray(np.array([1, 0]), like=framework)
+
+    numpy_data = np.array(0.62)
+
+    ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1)), qml.RZ(numpy_data, 0)]
+    m = [qml.state(), qml.expval(qml.Hermitian(M, 0))]
+    prep = [qml.QubitStateVector(state, 0)]
+
+    qs = qml.tape.QuantumScript(ops, m, prep)
+    new_qs = convert_to_numpy_parameters(qs)
+
+    # check ops that should be unaltered
+    assert new_qs[3] is qs[3]
+    assert new_qs[4] is qs[4]
+    assert new_qs.measurements[0] is qs.measurements[0]
+
+    for ind in (0, 1, 2, 6):
+        assert qml.equal(new_qs[ind], qs[ind], check_interface=False, check_trainability=False)
+        assert qml.math.get_interface(*new_qs[ind].data) == "numpy"
+
+
+@pytest.mark.autograd
+def test_preserves_trainable_params():
+    """Test that convert_to_numpy_parameters preserves the trainable parameters property."""
+    ops = [qml.RX(qml.numpy.array(2.0), 0), qml.RY(qml.numpy.array(3.0), 0)]
+    qs = qml.tape.QuantumScript(ops)
+    qs.trainable_params = {0}
+    output = convert_to_numpy_parameters(qs)
+    assert output.trainable_params == [0]
+
+
+@pytest.mark.autograd
+def test_unwraps_arithmetic_op():
+    """Test that the operator helper function can handle operator arithmetic objects."""
+    op1 = qml.s_prod(qml.numpy.array(2.0), qml.PauliX(0))
+    op2 = qml.s_prod(qml.numpy.array(3.0), qml.PauliY(0))
+
+    sum_op = qml.sum(op1, op2)
+
+    unwrapped_op = _convert_op_to_numpy_data(sum_op)
+    assert qml.math.get_interface(*unwrapped_op.data) == "numpy"
+    assert qml.math.get_interface(*unwrapped_op.data) == "numpy"
+
+
+@pytest.mark.autograd
+def test_unwraps_arithmetic_op_measurement():
+    """Test that the measurement helper function can handle operator arithmetic objects."""
+    op1 = qml.s_prod(qml.numpy.array(2.0), qml.PauliX(0))
+    op2 = qml.s_prod(qml.numpy.array(3.0), qml.PauliY(0))
+
+    sum_op = qml.sum(op1, op2)
+    m = qml.expval(sum_op)
+
+    unwrapped_m = _convert_measurement_to_numpy_data(m)
+    unwrapped_op = unwrapped_m.obs
+    assert qml.math.get_interface(*unwrapped_op.data) == "numpy"
+    assert qml.math.get_interface(*unwrapped_op.data) == "numpy"
+
+
+@pytest.mark.autograd
+def test_unwraps_tensor_observables():
+    """Test that the measurement helper function can set data on a tensor observable."""
+    mat = qml.numpy.eye(2)
+    obs = qml.operation.Tensor(qml.PauliZ(0), qml.Hermitian(mat, 1))
+    m = qml.expval(obs)
+
+    unwrapped_m = _convert_measurement_to_numpy_data(m)
+    assert qml.math.get_interface(*unwrapped_m.obs.data) == "numpy"