Add construct_batch and transform_program functions to workflow module

doc/releases/changelog-dev.md

@@ -116,6 +116,13 @@
 * Raise a more informative error when calling `adjoint_jacobian` with trainable state-prep operations.
   [(#5026)](https://github.com/PennyLaneAI/pennylane/pull/5026)
 
+* Adds `qml.workflow.transform_program` and `qml.workflow.construct_batch` to inspect the transform program and batch of tapes
+  at different stages.
+  [(#5084)](https://github.com/PennyLaneAI/pennylane/pull/5084)
+
+* Adds a `qml.transforms.core.expand_fn_transform` that converts a tape-to-tape function into a transform.
+  [(#5084)](https://github.com/PennyLaneAI/pennylane/pull/5084)
+
 * `CRX`, `CRY`, `CRZ`, `CROT`, and `ControlledPhaseShift` (i.e. `CPhaseShift`) now inherit from `ControlledOp`, giving them additional properties such as `control_wire` and `control_values`. Calling `qml.ctrl` on `RX`, `RY`, `RZ`, `Rot`, and `PhaseShift` with a single control wire will return gates of types `CRX`, `CRY`, etc. as opposed to a general `Controlled` operator.
   [(#5069)](https://github.com/PennyLaneAI/pennylane/pull/5069)
 

pennylane/_qubit_device.py

@@ -32,7 +32,6 @@
 
 import pennylane as qml
 from pennylane import Device, DeviceError
-from pennylane.workflow import set_shots
 from pennylane.math import multiply as qmlmul
 from pennylane.math import sum as qmlsum
 from pennylane.measurements import (
@@ -1036,7 +1035,7 @@ def classical_shadow(self, obs, circuit):
         n_snapshots = self.shots
         seed = obs.seed
 
-        with set_shots(self, shots=1):
+        with qml.workflow.set_shots(self, shots=1):
             # slow implementation but works for all devices
             n_qubits = len(wires)
             mapped_wires = np.array(self.map_wires(wires))

pennylane/devices/default_qubit.py

@@ -159,6 +159,16 @@ def adjoint_state_measurements(
     )
 
 
+def adjoint_ops(op: qml.operation.Operator) -> bool:
+    """Specify whether or not an Operator is supported by adjoint differentiation."""
+    return op.num_params == 0 or op.num_params == 1 and op.has_generator
+
+
+def adjoint_observables(obs: qml.operation.Operator) -> bool:
+    """Specifies whether or not an observable is compatible with adjoint differentiation on DefaultQubit."""
+    return obs.has_matrix
+
+
 def _add_adjoint_transforms(program: TransformProgram, device_vjp=False) -> None:
     """Private helper function for ``preprocess`` that adds the transforms specific
     for adjoint differentiation.
@@ -171,14 +181,6 @@ def _add_adjoint_transforms(program: TransformProgram, device_vjp=False) -> None
 
     """
 
-    def adjoint_ops(op: qml.operation.Operator) -> bool:
-        """Specify whether or not an Operator is supported by adjoint differentiation."""
-        return op.num_params == 0 or op.num_params == 1 and op.has_generator
-
-    def adjoint_observables(obs: qml.operation.Operator) -> bool:
-        """Specifies whether or not an observable is compatible with adjoint differentiation on DefaultQubit."""
-        return obs.has_matrix
-
     name = "adjoint + default.qubit"
     program.add_transform(no_sampling, name=name)
     program.add_transform(

pennylane/templates/subroutines/permute.py

@@ -138,6 +138,9 @@ def circuit():
 
     """
 
+    def __repr__(self):
+        return f"Permute({self.hyperparameters['permutation']}, wires={self.wires.tolist()})"
+
     num_wires = AnyWires
     grad_method = None
 

pennylane/transforms/core/__init__.py

@@ -14,6 +14,6 @@
 r"""This module contains the experimental transforms building blocks (core).
 """
 
-from .transform import transform
+from .transform import transform, expand_fn_transform
 from .transform_dispatcher import TransformDispatcher, TransformContainer, TransformError
 from .transform_program import TransformProgram