Make qml.QutritBasisStatePreparation JIT compatible
#6308
Conversation
|
Hello. You may have forgotten to update the changelog!
|
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #6308 +/- ##
==========================================
- Coverage 99.71% 99.70% -0.01%
==========================================
Files 444 444
Lines 42252 42261 +9
==========================================
+ Hits 42132 42138 +6
- Misses 120 123 +3 ☔ View full report in Codecov by Sentry. |
tests/templates/test_state_preparations/test_qutrit_basis_state_prep.py
Outdated
Show resolved
Hide resolved
| @@ -315,6 +315,44 @@ def matmul(tensor1, tensor2, like=None): | |||
| return ar.numpy.matmul(tensor1, tensor2, like=like) | |||
|
|
|||
|
|
|||
| @multi_dispatch(argnum=[0, 1]) | |||
There was a problem hiding this comment.
With each new function added, we should include a test in test/math/test_multi_dispactch
| tshift = qml.math.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]]) | ||
| for wire, state in zip(wires, basis_state): | ||
| for _ in range(0, state): | ||
| op_list.append(qml.TShift(wire)) | ||
| mat = qml.math.matrix_power(tshift, state) | ||
| op = qml.ops.QutritUnitary(mat, wires=wire) | ||
| op_list.append(op) | ||
|
|
There was a problem hiding this comment.
to avoid QubitUnitary (since it is not hardware compatible), you could use this:
qml.pow(qml.TShift(wires = wire),state)
qml.pow returns an operator, so the idea is the same without access the matrix.
Note that if this works, you may not need to add matrix power in the dispach.py file
There was a problem hiding this comment.
The issue is that the pow method for qml.TShift is not jittable. It uses the pow method from Operator which has control flow to check if the exponent is 0 or 1.
Co-authored-by: Guillermo Alonso-Linaje <65235481+KetpuntoG@users.noreply.github.com>
There was a problem hiding this comment.
Thanks @willjmax! Some minor queries, before I dig in deeper.
| for wire, state in zip(wires, basis_state): | ||
| for _ in range(0, state): | ||
| op_list.append(qml.TShift(wire)) | ||
| mat = qml.math.matrix_power(tshift, state) |
There was a problem hiding this comment.
We cannot use qml.math.linalg.matrix_power? 🤔
| @@ -112,7 +113,10 @@ def compute_decomposition(basis_state, wires): # pylint: disable=arguments-diff | |||
| """ | |||
|
|
|||
| op_list = [] | |||
| tshift = qml.math.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]]) | |||
There was a problem hiding this comment.
Does this scale worse or better than the default case? If worse, then should we limit this only to jit case (and use the default behaviour otherwise)?
This PR makes
qml.QutritBasisStatePreparationJIT compatible. Previously, the template used theqml.TShiftoperator in the decomposition. However, this approach requires non-jittable control flow on the input. To make the template jittable the decomposition is changed to return aqml.ops.QutritUnitaryrepresenting aTShiftoperator raised to some power. This change requires addingmatrix_powerandeighto the multi dispatch. This PR addresses sc-70863.