Repository: PennyLaneAI/pennylane
Pull request #5854 (2024-06-13 21:04:11)

Fix `level=user` bug in `get_transform_program`
**Context:**
When requesting `level=user` from `get_transform_program` with a final transform, only the expand function of that transform was included without the actual transform. 

**Description of the Change:**
The final transform is now added if it exists.

[[sc-66171](https://app.shortcut.com/xanaduai/story/66171)]



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://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5854?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.67%. Comparing base [(`cdc36bc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/cdc36bc6b8e8ea55b2ff8664fd14da5a3078f3d9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`e06ee24`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e06ee2422835f58137e3e1ec6ff52e91dbf19825?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5854      +/-   ##
==========================================
- Coverage   99.67%   99.67%   -0.01%     
==========================================
  Files         420      420              
  Lines       40170    39879     -291     
==========================================
- Hits        40041    39749     -292     
- Misses        129      130       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5854?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/pennylane/workflow/construct_batch.py b/pennylane/workflow/construct_batch.py
index bc537df6b..717fa9d8e 100644
--- a/pennylane/workflow/construct_batch.py
+++ b/pennylane/workflow/construct_batch.py
@@ -92,9 +92,10 @@ def get_transform_program(qnode: "QNode", level=None) -> "qml.transforms.core.Tr
             :width: 800px
             :target: javascript:void(0);
 
-        where ``transform1`` is first applied to the ``QNode`` followed by ``transform2``.  First user transforms are run on the tapes,
-        followed by the gradient expansion, followed by the device expansion.  "Final" transforms, like ``param_shift`` and ``metric_tensor``,
-        always occur at the end of the program.
+        where ``transform1`` is first applied to the ``QNode`` followed by ``transform2``.  First, user transforms are run on the tapes,
+        followed by the gradient expansion, followed by the device expansion. "Final" transforms, like ``param_shift`` and ``metric_tensor``,
+        always occur at the end of the program, despite being part of user transforms. Note that when requesting a level by name
+        (e.g. "gradient" or "device"), the preceding levels would be applied as well.
 
         .. code-block:: python
 
@@ -114,26 +115,27 @@ def get_transform_program(qnode: "QNode", level=None) -> "qml.transforms.core.Tr
         _expand_transform_param_shift, validate_device_wires, defer_measurements,
         decompose, validate_measurements, validate_observables, metric_tensor)
 
-        The ``"user"`` transforms are the ones manually applied to the qnode, :class:`~.cancel_inverses` and
-        :class:`~.merge_rotations`.
+        The ``"user"`` transforms are the ones manually applied to the qnode, :func:`~.cancel_inverses`,
+        :func:`~.merge_rotations` and :func:`~.metric_tensor`.
 
         >>> qml.workflow.get_transform_program(circuit, level="user")
-        TransformProgram(cancel_inverses, merge_rotations)
+        TransformProgram(cancel_inverses, merge_rotations, _expand_metric_tensor, metric_tensor)
 
-        The ``_expand_transform_param_shift`` is the ``"gradient"`` transform.  This expands all trainable
-        operations to a state where the parameter shift transform can operate on them. For example, it will decompose
-        any parametrized templates into operators that have generators.
+        The ``_expand_transform_param_shift`` is the ``"gradient"`` transform.
+        This expands all trainable operations to a state where the parameter shift transform can operate on them. For example,
+        it will decompose any parametrized templates into operators that have generators. Note how ``metric_tensor`` is still
+        present at the very end of resulting program.
 
         >>> qml.workflow.get_transform_program(circuit, level="gradient")
-        TransformProgram(cancel_inverses, merge_rotations, _expand_transform_param_shift)
+        TransformProgram(cancel_inverses, merge_rotations, _expand_metric_tensor, _expand_transform_param_shift, metric_tensor)
 
-        ``"device"`` includes all transforms except for a ``"final"`` transform, if it exists.  This usually
+        ``"device"`` is equivalent to ``level=None`` and includes all transforms. Semantically, this usually
         corresponds to the circuits that will be sent to the device to execute.
 
         >>> qml.workflow.get_transform_program(circuit, level="device")
         TransformProgram(cancel_inverses, merge_rotations, _expand_transform_param_shift,
         validate_device_wires, defer_measurements, decompose, validate_measurements,
-        validate_observables)
+        validate_observables, metric_tensor)
 
         ``"top"`` and ``0`` both return empty transform programs.
 
@@ -158,6 +160,17 @@ def get_transform_program(qnode: "QNode", level=None) -> "qml.transforms.core.Tr
         decompose, defer_measurements, validate_device_wires, _expand_transform_param_shift,
         _expand_metric_tensor, merge_rotations, cancel_inverses)
 
+        You can get creative and pick a single category of transforms as follows, excluding
+        any preceding transforms (and the final transform if it exists):
+
+        >>> user_prog = qml.workflow.get_transform_program(circuit, level="user")
+        >>> grad_prog = qml.workflow.get_transform_program(circuit, level="gradient")
+        >>> dev_prog = qml.workflow.get_transform_program(circuit, level="device")
+        >>> grad_prog[len(user_prog) - 1 : -1]
+        TransformProgram(_expand_transform_param_shift)
+        >>> dev_prog[len(grad_prog) - 1 : -1]
+        TransformProgram(validate_device_wires, mid_circuit_measurements, decompose, validate_measurements, validate_observables)
+
     """
     full_transform_program = _get_full_transform_program(qnode)
 
@@ -166,24 +179,33 @@ def get_transform_program(qnode: "QNode", level=None) -> "qml.transforms.core.Tr
         # final transform is placed after device transforms
         num_user -= 1
 
+    readd_final_transform = False
+
     if level == "device":
-        level = -1 if full_transform_program.has_final_transform else None
+        level = None
     elif level == "top":
         level = 0
     elif level == "user":
+        readd_final_transform = True
         level = num_user
     elif level == "gradient":
-        if getattr(qnode.gradient_fn, "expand_transform", False):
-            level = slice(0, num_user + 1)
-        else:
-            level = slice(0, num_user)
+        readd_final_transform = True
+
+        level = num_user + 1 if getattr(qnode.gradient_fn, "expand_transform", False) else num_user
     elif isinstance(level, str):
         raise ValueError(
             f"level {level} not recognized. Acceptable strings are 'device', 'top', 'user', and 'gradient'."
         )
+
     if level is None or isinstance(level, int):
         level = slice(0, level)
-    return full_transform_program[level]
+
+    resolved_program = full_transform_program[level]
+
+    if qnode.transform_program.has_final_transform and readd_final_transform:
+        resolved_program += qnode.transform_program[-1:]
+
+    return resolved_program
 
 
 def construct_batch(qnode: QNode, level: Union[None, str, int, slice] = "user") -> Callable:


Test patch:
diff --git a/tests/workflow/test_construct_batch.py b/tests/workflow/test_construct_batch.py
index 44e345e55..98c7f30f3 100644
--- a/tests/workflow/test_construct_batch.py
+++ b/tests/workflow/test_construct_batch.py
@@ -186,23 +186,27 @@ class TestTransformProgramGetter:
             return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(0))
 
         user_program = get_transform_program(circuit, level="user")
-        assert len(user_program) == 2
+        assert len(user_program) == 3
         assert user_program[0].transform == qml.compile.transform
         assert user_program[1].transform == qml.metric_tensor.expand_transform
+        assert user_program[2].transform == qml.metric_tensor.transform
 
         grad_program = get_transform_program(circuit, level="gradient")
-        assert len(grad_program) == 3
+        assert len(grad_program) == 4
         assert grad_program[0].transform == qml.compile.transform
         assert grad_program[1].transform == qml.metric_tensor.expand_transform
         assert grad_program[2].transform == qml.gradients.param_shift.expand_transform
+        assert grad_program[3].transform == qml.metric_tensor.transform
 
         dev_program = get_transform_program(circuit, level="device")
         config = qml.devices.ExecutionConfig(interface=getattr(circuit, "interface", None))
-        assert len(dev_program) == 3 + len(circuit.device.preprocess(config)[0])  # currently 8
-        assert qml.metric_tensor not in dev_program
+        assert len(dev_program) == 4 + len(circuit.device.preprocess(config)[0])  # currently 8
+        assert dev_program[-1].transform == qml.metric_tensor.transform
 
-        full = get_transform_program(circuit)
-        assert full[-1].transform == qml.metric_tensor.transform
+        full_program = get_transform_program(circuit)
+        assert full_program[-1].transform == qml.metric_tensor.transform
+
+        assert dev_program == full_program
 
 
 @qml.transforms.merge_rotations