Incorporate level keyword in draw and draw_mpl (#5855)

**Context:**
Currently, `draw()` and `draw_mpl()` can only be requested after
applying the full transform program, with the exception of the stages
provided through expansion_strategy.

**Description of the Change:**
Using the new `level` argument from construct_batch, the functions are
adapted to make use of the argument as well, which allows for more
flexible requests and ability to pinpoint where exactly in the transform
program to draw a circuit. As done before, the new functionality works
with transforms that split the tape (only in the case for `draw`. For
`draw_mpl`, a warning is raised and only the first tape is plotted).

**Benefits:**
Better plotting UX.

**Note for Reviewers:**
Minor bugs in `construct_batch` have been discovered during work on this
PR, and so expect minor fixes to tests relating to `specs`.

[[sc-53735](https://app.shortcut.com/xanaduai/story/53735)] Supersedes
#5139

---------

Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com>

doc/_static/draw_mpl/draw_mpl_examples.py

@@ -102,7 +102,7 @@ def rcparams(circuit):
 
 def use_style(circuit):
 
- fig, ax = qml.draw_mpl(circuit, style='sketch')(1.2345, 1.2345)
+ fig, ax = qml.draw_mpl(circuit, style="sketch")(1.2345, 1.2345)
 
  plt.savefig(folder / "sketch_style.png")
  plt.close()
@@ -128,6 +128,26 @@ def circuit():
  plt.close()
 
 
+@qml.transforms.merge_rotations
+@qml.transforms.cancel_inverses
+@qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift")
+def _levels_circ():
+ qml.RandomLayers([[1.0, 20]], wires=(0, 1))
+ qml.Permute([2, 1, 0], wires=(0, 1, 2))
+ qml.PauliX(0)
+ qml.PauliX(0)
+ qml.RX(0.1, wires=0)
+ qml.RX(-0.1, wires=0)
+ return qml.expval(qml.PauliX(0))
+
+
+def levels():
+ for level in ("top", "user", None, slice(1, 2)):
+ draw_mpl(_levels_circ, level=level)()
+ plt.savefig(folder / f"level_{str(level).split('(')[0].lower()}.png")
+ plt.close
+
+
 if __name__ == "__main__":
 
  dev = qml.device("lightning.qubit", wires=(0, 1, 2, 3))
@@ -151,3 +171,4 @@ def circuit(x, z):
  rcparams(circuit)
  wires_labels(circuit)
  mid_measure()
+ levels()

doc/releases/changelog-dev.md

@@ -307,6 +307,9 @@
 * `qml.qchem.molecular_dipole` function is added for calculating the dipole operator using "dhf" and "openfermion" backends.
  [(#5764)](https://github.com/PennyLaneAI/pennylane/pull/5764)
 
+* Circuits can now be plotted at any specific point of the transform program through the `level` keyword argument in `draw()` and `draw_mpl()`.
+ [(#5855)](https://github.com/PennyLaneAI/pennylane/pull/5855)
+
 * Transforms applied to callables now use `functools.wraps` to preserve the docstring and call signature of the original function.
  [(#5857)](https://github.com/PennyLaneAI/pennylane/pull/5857)
 

pennylane/qnn/torch.py

@@ -458,7 +458,7 @@ def construct(self, args, kwargs):
  x = args[0]
  kwargs = {
  self.input_arg: x,
- **{arg: weight.data.to(x) for arg, weight in self.qnode_weights.items()},
+ **{arg: weight.to(x) for arg, weight in self.qnode_weights.items()},
  }
  self.qnode.construct((), kwargs)
 

pennylane/resource/specs.py

@@ -121,71 +121,71 @@ def circuit(x, add_ry=True):
  'gradient_fn': 'pennylane.gradients.parameter_shift.param_shift',
  'num_gradient_executions': 2}
 
- Here you can see how the number of gates and their types change as we apply different amounts of transforms
- through the ``level`` argument:
+ .. details::
+  :title: Usage Details
 
- .. code-block:: python3
-
- @qml.transforms.merge_rotations
- @qml.transforms.undo_swaps
- @qml.transforms.cancel_inverses
- @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift", shifts=np.pi / 4)
- def circuit(x):
- qml.RandomLayers(qml.numpy.array([[1.0, 2.0]]), wires=(0, 1))
- qml.RX(x, wires=0)
- qml.RX(-x, wires=0)
- qml.SWAP((0, 1))
- qml.X(0)
- qml.X(0)
- return qml.expval(qml.X(0) + qml.Y(1))
+ Here you can see how the number of gates and their types change as we apply different amounts of transforms
+ through the ``level`` argument:
 
- return circuit
+ .. code-block:: python3
 
- First, we can check the resource information of the ``QNode`` without any modifications. Note that ``level=top`` would
- return the same results:
+ @qml.transforms.merge_rotations
+ @qml.transforms.undo_swaps
+ @qml.transforms.cancel_inverses
+ @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift", shifts=np.pi / 4)
+ def circuit(x):
+ qml.RandomLayers(qml.numpy.array([[1.0, 2.0]]), wires=(0, 1))
+ qml.RX(x, wires=0)
+ qml.RX(-x, wires=0)
+ qml.SWAP((0, 1))
+ qml.X(0)
+ qml.X(0)
+ return qml.expval(qml.X(0) + qml.Y(1))
 
- >>> qml.specs(circuit, level=0)(0.1)["resources"]
- Resources(num_wires=2, num_gates=6, gate_types=defaultdict(<class 'int'>, {'RandomLayers': 1, 'RX': 2, 'SWAP': 1, 'PauliX': 2}),
- gate_sizes=defaultdict(<class 'int'>, {2: 2, 1: 4}), depth=6, shots=Shots(total_shots=None, shot_vector=()))
+ First, we can check the resource information of the ``QNode`` without any modifications. Note that ``level=top`` would
+ return the same results:
 
- We then check the resources after applying all transforms:
+ >>> qml.specs(circuit, level=0)(0.1)["resources"]
+ Resources(num_wires=2, num_gates=6, gate_types=defaultdict(<class 'int'>, {'RandomLayers': 1, 'RX': 2, 'SWAP': 1, 'PauliX': 2}),
+ gate_sizes=defaultdict(<class 'int'>, {2: 2, 1: 4}), depth=6, shots=Shots(total_shots=None, shot_vector=()))
 
- >>> qml.specs(circuit, level=None)(0.1)["resources"]
- Resources(num_wires=2, num_gates=2, gate_types=defaultdict(<class 'int'>, {'RY': 1, 'RX': 1}),
- gate_sizes=defaultdict(<class 'int'>, {1: 2}), depth=1, shots=Shots(total_shots=None, shot_vector=()))
+ We then check the resources after applying all transforms:
 
- We can also notice that ``SWAP`` and ``PauliX`` are not present in the circuit if we set ``level=2``:
+ >>> qml.specs(circuit, level=None)(0.1)["resources"]
+ Resources(num_wires=2, num_gates=2, gate_types=defaultdict(<class 'int'>, {'RY': 1, 'RX': 1}),
+ gate_sizes=defaultdict(<class 'int'>, {1: 2}), depth=1, shots=Shots(total_shots=None, shot_vector=()))
 
- >>> qml.specs(circuit, level=2)(0.1)["resources"]
- Resources(num_wires=2, num_gates=3, gate_types=defaultdict(<class 'int'>, {'RandomLayers': 1, 'RX': 2}),
- gate_sizes=defaultdict(<class 'int'>, {2: 1, 1: 2}), depth=3, shots=Shots(total_shots=None, shot_vector=()))
+ We can also notice that ``SWAP`` and ``PauliX`` are not present in the circuit if we set ``level=2``:
 
- If we attempt to only apply the ``merge_rotations`` transform, we would end with only one trainable object, which is in ``RandomLayers``:
+ >>> qml.specs(circuit, level=2)(0.1)["resources"]
+ Resources(num_wires=2, num_gates=3, gate_types=defaultdict(<class 'int'>, {'RandomLayers': 1, 'RX': 2}),
+ gate_sizes=defaultdict(<class 'int'>, {2: 1, 1: 2}), depth=3, shots=Shots(total_shots=None, shot_vector=()))
 
- >>> qml.specs(circuit, level=slice(2, 3))(0.1)["num_trainable_params"]
- 1
+ If we attempt to only apply the ``merge_rotations`` transform, we would end with only one trainable object, which is in ``RandomLayers``:
 
- However, if we apply all transforms, ``RandomLayers`` would be decomposed to an ``RY`` and an ``RX``, giving us two trainable objects:
+ >>> qml.specs(circuit, level=slice(2, 3))(0.1)["num_trainable_params"]
+ 1
 
- >>> qml.specs(circuit, level=None)(0.1)["num_trainable_params"]
- 2
+ However, if we apply all transforms, ``RandomLayers`` would be decomposed to an ``RY`` and an ``RX``, giving us two trainable objects:
 
- If a ``QNode`` with a tape-splitting transform is supplied to the function, with the transform included in the desired transforms, a dictionary
- would be returned for each resulting tapes:
+  >>> qml.specs(circuit, level=None)(0.1)["num_trainable_params"]
+  2
 
- .. code-block:: python3
+ If a ``QNode`` with a tape-splitting transform is supplied to the function, with the transform included in the desired transforms, a dictionary
+ would be returned for each resulting tapes:
 
- H = qml.Hamiltonian([0.2, -0.543], [qml.X(0) @ qml.Z(1), qml.Z(0) @ qml.Y(2)])
+ .. code-block:: python3
 
+ H = qml.Hamiltonian([0.2, -0.543], [qml.X(0) @ qml.Z(1), qml.Z(0) @ qml.Y(2)])
 
- @qml.transforms.hamiltonian_expand
- @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift", shifts=np.pi / 4)
- def circuit():
- qml.RandomLayers(qml.numpy.array([[1.0, 2.0]]), wires=(0, 1))
- return qml.expval(H)
+  @qml.transforms.hamiltonian_expand
+  @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift", shifts=np.pi / 4)
+  def circuit():
+  qml.RandomLayers(qml.numpy.array([[1.0, 2.0]]), wires=(0, 1))
+  return qml.expval(H)
 
- >>> len(qml.specs(circuit, level="user")())
- 2
+  >>> len(qml.specs(circuit, level="user")())
+  2
  """
 
  specs_level = _determine_spec_level(kwargs, qnode)

pennylane/workflow/construct_batch.py

@@ -15,6 +15,7 @@
 
 """
 import inspect
+from contextlib import nullcontext
 from functools import wraps
 from typing import Callable, Literal, Optional, Tuple, Union
 
@@ -54,16 +55,20 @@ def wrapped_expand_fn(tape, *args, **kwargs):
 
 def _get_full_transform_program(qnode: QNode) -> "qml.transforms.core.TransformProgram":
  program = qml.transforms.core.TransformProgram(qnode.transform_program)
+
  if getattr(qnode.gradient_fn, "expand_transform", False):
  program.add_transform(
  qml.transform(qnode.gradient_fn.expand_transform),
  **qnode.gradient_kwargs,
  )
+
  if isinstance(qnode.device, qml.devices.Device):
  config = _make_execution_config(qnode, qnode.gradient_fn)
  return program + qnode.device.preprocess(config)[0]
+
  program.add_transform(qml.transform(qnode.device.batch_transform))
  program.add_transform(expand_fn_transform(qnode.device.expand_fn))
+
  return program
 
 
@@ -316,28 +321,37 @@ def batch_constructor(*args, **kwargs) -> Tuple[Tuple["qml.tape.QuantumTape", Ca
  else:
  shots = kwargs.pop("shots", _get_device_shots(qnode.device))
 
+ context_fn = nullcontext
+
  if isinstance(qnode, qml.qnn.KerasLayer):
  # pylint: disable=import-outside-toplevel
  import tensorflow as tf
 
- with tf.GradientTape() as tape:
- tape.watch(list(qnode.qnode_weights.values()))
-
- kwargs = {
- **{k: 1.0 * w for k, w in qnode.qnode_weights.items()},
- **kwargs,
- }
+ context_fn = tf.GradientTape
 
  if isinstance(qnode, qml.qnn.TorchLayer):
  x = args[0]
  kwargs = {
- **{arg: weight.data.to(x) for arg, weight in qnode.qnode_weights.items()},
+ **{arg: weight.to(x) for arg, weight in qnode.qnode_weights.items()},
  }
 
- initial_tape = qml.tape.make_qscript(qnode.func, shots=shots)(*args, **kwargs)
+ with context_fn() as cntxt:
+ # If TF tape, use the watch function
+ if hasattr(cntxt, "watch"):
+ cntxt.watch(list(qnode.qnode_weights.values()))
+
+ kwargs = {
+ **{k: 1.0 * w for k, w in qnode.qnode_weights.items()},
+ **kwargs,
+ }
+
+ initial_tape = qml.tape.make_qscript(qnode.func, shots=shots)(*args, **kwargs)
+ params = initial_tape.get_parameters(trainable_only=False)
+ initial_tape.trainable_params = qml.math.get_trainable_indices(params)
 
  qnode._update_gradient_fn(tape=initial_tape)
  program = get_transform_program(qnode, level=level)
+
  return program((initial_tape,))
 
  return batch_constructor