Ensure that `MottonenStatePreparation` is fully differentiable

[mariaschuld]

This issue has been tagged for contributions during unitaryHACK

Issue description

At the moment, the MottonenStatePreparation template that prepares a quantum state whose Dirac vector encodes some specified input vector, is not fully differentiable. This is why the test cases for differentiability/interfaces are currently reduced to a trivial example.

• Expected behavior: (What you expect to happen)

The template should be fully differentiable with respect to state_vector in all interfaces (tf, torch, autograd/numpy, jax).

• Actual behavior: (What actually happens)

At the moment, differentiation in some interfaces does not work. This is likely due to the complicated input processing pipelines in the template, which were rewritten in PennyLane's math library to be interface-agnostic, but which do not seem to work 100% yet.

A detailed investigation about what cases work and which ones don't may be part of the solution to this issue.

Additional information

A minimum goal could be to copy the differentiation tests from another template (like this class of tests example) and to make it pass. To come up with a circuit that one can compare the template against, one would have to figure out the decomposition (and how angles of the gates depend on the input vector to the template) for a very simple case, which may be a bit tricky.

MottonenStatePreparation is called by AmplitudeEmbedding and QubitStateVector if the device needs a decomposition. By extension, these two templates/operations should become differentiable on more devices.

[josh146]

The following discussion forum post provides a bug fix to get differentiability working with PyTorch: https://discuss.pennylane.ai/t/hybrid-network-not-differentiating/1079/3

[josh146]

In particular, here is the relevant git diff:

diff --git a/pennylane/templates/state_preparations/mottonen.py b/pennylane/templates/state_preparations/mottonen.py
index 860eeb80..a7a83544 100644
--- a/pennylane/templates/state_preparations/mottonen.py
+++ b/pennylane/templates/state_preparations/mottonen.py
@@ -203,6 +203,9 @@ def _get_alpha_y(a, n, k):
     with np.errstate(divide="ignore", invalid="ignore"):
         division = numerator / denominator

+    division = qml.math.cast(division, np.float64)
+    denominator = qml.math.cast(denominator, np.float64)
+
     division = qml.math.where(denominator != 0.0, division, 0.0)

     return 2 * qml.math.arcsin(qml.math.sqrt(division))

[antalszava]

The fix was incorporated in https://github.com/PennyLaneAI/pennylane/pull/1400. One test that reproduced the bug was added, though we could extend that with the entire class of tests as suggested in the original description.