Remove the stacking (and transposing) behaviour of qml.jacobian (#2059

)

* main commit

* black

* changelog

* lint

* docstrings

* correct changelog

* black

* remove prints in tests

* add comment regarding QNGOptimizer and multiple arugments

* black

* add accidentally deleted print

* inline code blocks

* Apply suggestions from code review

Co-authored-by: Josh Izaac <josh146@gmail.com>

* basic example qml.jacobian

* qng todo comment

* intermediate

* tests

* Apply suggestions from code review

Co-authored-by: David Ittah <dime10@users.noreply.github.com>

* Introduce `qml.math.safe_squeeze` and remove excessive squeezes from `gradient_transform` (#2080)

* safe_squeeze+tests

* gradient trafo

* include previously skipped tests

* black

* changelog

* Apply suggestions from code review

Co-authored-by: Josh Izaac <josh146@gmail.com>

Co-authored-by: Josh Izaac <josh146@gmail.com>

* code review: doc string and old note

* table column size

* exception generalization to accomodate for TF exceptions in safe_squeeze

* concrete exceptions in safe_squeeze

* column widths

* tf error in safe_squeeze

* rotosolve linting gone wrong, included here

* remove unused scenario in QNG

* extend comment; remove condition for linting purpose

Co-authored-by: Josh Izaac <josh146@gmail.com>
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: antalszava <antalszava@gmail.com>

doc/releases/changelog-dev.md

@@ -216,6 +216,18 @@
 
 <h3>Improvements</h3>
 
+* The function `qml.math.safe_squeeze` is introduced and `gradient_transform` allows
+  for QNode argument axes of size `1`.
+  [(#2080)](https://github.com/PennyLaneAI/pennylane/pull/2080)
+
+  `qml.math.safe_squeeze` wraps `qml.math.squeeze`, with slight modifications:
+
+  - When provided the `axis` keyword argument, axes that do not have size `1` will be
+    ignored, instead of raising an error.
+
+  - The keyword argument `exclude_axis` allows to explicitly exclude axes from the
+    squeezing.
+
 * The `adjoint` transform now raises and error whenever the object it is applied to
   is not callable.
   [(#2060)](https://github.com/PennyLaneAI/pennylane/pull/2060)
@@ -279,6 +291,35 @@
 
 <h3>Breaking changes</h3>
 
+* `qml.metric_tensor`, `qml.adjoint_metric_tensor` and `qml.transforms.classical_jacobian`
+  now follow a different convention regarding their output shape when being used
+  with the Autograd interface
+  [(#2059)](https://github.com/PennyLaneAI/pennylane/pull/2059)
+
+  See the previous entry for details. This breaking change immediately follows from
+  the change in `qml.jacobian` whenever `hybrid=True` is used in the above methods.
+
+* `qml.jacobian` now follows a different convention regarding its output shape.
+  [(#2059)](https://github.com/PennyLaneAI/pennylane/pull/2059)
+
+  Previously, `qml.jacobian` would attempt to stack the Jacobian for multiple
+  QNode arguments, which succeeded whenever the arguments have the same shape.
+  In this case, the stacked Jacobian would also be transposed, leading to the 
+  output shape `(*reverse_QNode_args_shape, *reverse_output_shape, num_QNode_args)`
+
+  If no stacking and transposing occurs, the output shape instead is a `tuple`
+  where each entry corresponds to one QNode argument and has the shape
+  `(*output_shape, *QNode_arg_shape)`.
+
+  This breaking change alters the behaviour in the first case and removes the attempt
+  to stack and transpose, so that the output always has the shape of the second
+  type.
+
+  Note that the behaviour is unchanged --- that is, the Jacobian tuple is unpacked into
+  a single Jacobian --- if `argnum=None` and there is only one QNode argument
+  with respect to which the differentiation takes place, or if an integer
+  is provided as `argnum`.
+
 * The behaviour of `RotosolveOptimizer` has been changed regarding
   its keyword arguments.
   [(#2081)](https://github.com/PennyLaneAI/pennylane/pull/2081)
@@ -301,9 +342,13 @@
   For more details, see the
   [RotosolveOptimizer documentation](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.RotosolveOptimizer.html).
 
-
 <h3>Bug fixes</h3>
 
+* Fixes a bug in `gradient_transform` where the hybrid differentiation
+  of circuits with a single parametrized gate failed and QNode argument
+  axes of size `1` where removed from the output gradient.
+  [(#2080)](https://github.com/PennyLaneAI/pennylane/pull/2080)
+
 * The available `diff_method` options for QNodes has been corrected in both the
   error messages and the documentation.
   [(#2078)](https://github.com/PennyLaneAI/pennylane/pull/2078)

pennylane/_grad.py

@@ -24,8 +24,6 @@
 from autograd.extend import vspace
 from autograd.wrap_util import unary_to_nary
 
-from pennylane import numpy as np
-
 make_vjp = unary_to_nary(_make_vjp)
 
 
@@ -158,43 +156,147 @@ def jacobian(func, argnum=None):
             must consist of a single NumPy array (if classical) or a tuple of
             expectation values (if a quantum node)
         argnum (int or Sequence[int]): Which argument to take the gradient
-            with respect to. If a sequence is given, the Jacobian matrix
-            corresponding to all input elements and all output elements is returned.
+            with respect to. If a sequence is given, the Jacobian corresponding
+            to all marked inputs and all output elements is returned.
 
     Returns:
         function: the function that returns the Jacobian of the input
         function with respect to the arguments in argnum
-    """
-    # pylint: disable=no-value-for-parameter
 
-    if argnum is not None:
-        # for backwards compatibility with existing code
-        # that manually specifies argnum
-        if isinstance(argnum, int):
-            return _jacobian(func, argnum)
+    For ``argnum=None``, the trainable arguments are inferred dynamically from the arguments
+    passed to the function. The returned function takes the same arguments as the original
+    function and outputs a ``tuple``. The ``i`` th entry of the ``tuple`` has shape
+    ``(*output shape, *shape of args[argnum[i]])``.
 
-        return lambda *args, **kwargs: np.stack(
-            [_jacobian(func, arg)(*args, **kwargs) for arg in argnum]
-        ).T
+    If a single trainable argument is inferred, or if a single integer
+    is provided as ``argnum``, the tuple is unpacked and its only entry is returned instead.
 
-    def _jacobian_function(*args, **kwargs):
-        """Inspect the arguments for differentiability, and
-        compute the autograd gradient function with required argnums
-        dynamically.
+    **Example**
 
-        This wrapper function is returned to the user instead of autograd.jacobian,
-        so that we can take into account cases where the user computes the
-        jacobian function once, but then calls it with arguments that change
-        in differentiability.
-        """
+    Consider the QNode
+
+    .. code-block::
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[0, 0, 0], wires=0)
+            qml.RY(weights[0, 0, 1], wires=1)
+            qml.RZ(weights[1, 0, 2], wires=0)
+            return tuple(qml.expval(qml.PauliZ(w)) for w in dev.wires)
+
+        weights = np.array(
+            [[[0.2, 0.9, -1.4]], [[0.5, 0.2, 0.1]]], requires_grad=True
+        )
+
+    It has a single array-valued QNode argument with shape ``(2, 1, 3)`` and outputs
+    a tuple of two expectation values. Therefore, the Jacobian of this QNode
+    will be a single array with shape ``(2, 2, 1, 3)``:
+
+    >>> qml.jacobian(circuit)(weights).shape
+    (2, 2, 1, 3)
+
+    On the other hand, consider the following QNode for the same circuit
+    structure:
+
+    .. code-block::
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, y, z):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.RZ(z, wires=0)
+            return tuple(qml.expval(qml.PauliZ(w)) for w in dev.wires)
+
+        x = np.array(0.2, requires_grad=True)
+        y = np.array(0.9, requires_grad=True)
+        z = np.array(-1.4, requires_grad=True)
+
+    It has three scalar QNode arguments and outputs a tuple of two expectation
+    values. Consequently, its Jacobian will be a three-tuple of arrays with the
+    shape ``(2,)``:
+
+    >>> jac = qml.jacobian(circuit)(x, y, z)
+    >>> type(jac)
+    tuple
+    >>> for sub_jac in jac:
+    ...     print(sub_jac.shape)
+    (2,)
+    (2,)
+    (2,)
+
+    For a more advanced setting of QNode arguments, consider the QNode
+
+    .. code-block::
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x[0], wires=0)
+            qml.RY(y[0, 3], wires=1)
+            qml.RX(x[1], wires=2)
+            return [qml.expval(qml.PauliZ(w)) for w in [0, 1, 2]]
+
+        x = np.array([0.1, 0.5], requires_grad=True)
+        y = np.array([[-0.3, 1.2, 0.1, 0.9], [-0.2, -3.1, 0.5, -0.7]], requires_grad=True)
+
+    If we do not provide ``argnum``, ``qml.jacobian`` will correctly identify both,
+    ``x`` and ``y``, as trainable function arguments:
+
+    >>> jac = qml.jacobian(circuit)(x, y)
+    >>> print(type(jac), len(jac))
+    <class 'tuple'> 2
+    >>> qml.math.shape(jac[0])
+    (3, 2)
+    >>> qml.math.shape(jac[1])
+    (3, 2, 4)
+
+    As we can see, there are two entries in the output, one Jacobian for each
+    QNode argument. The shape ``(3, 2)`` of the first Jacobian is the combination
+    of the QNode output shape (``(3,)``) and the shape of ``x`` (``(2,)``).
+    Similarily, the shape ``(2, 4)`` of ``y`` leads to a Jacobian shape ``(3, 2, 4)``.
+
+    Instead we may choose the output to contain only one of the two
+    entries by providing an iterable as ``argnum``:
+
+    >>> jac = qml.jacobian(circuit, argnum=[1])(x, y)
+    >>> print(type(jac), len(jac))
+    <class 'tuple'> 1
+    >>> qml.math.shape(jac)
+    (1, 3, 2, 4)
+
+    Here we included the size of the tuple in the shape analysis, corresponding to the
+    first dimension of size ``1``.
+
+    Finally, we may want to receive the single entry above directly, not as a tuple
+    with a single entry. This is done by providing a single integer as ``argnum``
+
+    >>> jac = qml.jacobian(circuit, argnum=1)(x, y)
+    >>> print(type(jac), len(jac))
+    <class 'numpy.ndarray'> 3
+    >>> qml.math.shape(jac)
+    (3, 2, 4)
+
+    As expected, the tuple was unpacked and we directly received the Jacobian of the
+    QNode with respect to ``y``.
+    """
+    # pylint: disable=no-value-for-parameter
+
+    def _get_argnum(args):
+        """Inspect the arguments for differentiability and return the
+        corresponding indices."""
         argnum = []
 
         for idx, arg in enumerate(args):
 
             trainable = getattr(arg, "requires_grad", None)
-            array_box = isinstance(arg, ArrayBox)
+            is_array_box = isinstance(arg, ArrayBox)
 
-            if trainable is None and not array_box:
+            if trainable is None and not is_array_box:
 
                 warnings.warn(
                     "Starting with PennyLane v0.21.0, when using Autograd, inputs "
@@ -210,20 +312,29 @@ def _jacobian_function(*args, **kwargs):
             if trainable:
                 argnum.append(idx)
 
-        if not argnum:
-            return tuple()
-
-        if len(argnum) == 1:
-            return _jacobian(func, argnum[0])(*args, **kwargs)
+        return argnum
 
-        jacobians = [_jacobian(func, arg)(*args, **kwargs) for arg in argnum]
+    def _jacobian_function(*args, **kwargs):
+        """Compute the autograd Jacobian.
 
-        try:
-            return np.stack(jacobians).T
-        except ValueError:
-            # The Jacobian of each argument is a different shape and cannot
-            # be stacked; simply return the tuple of argument Jacobians.
-            return tuple(jacobians)
+        This wrapper function is returned to the user instead of autograd.jacobian,
+        so that we can take into account cases where the user computes the
+        jacobian function once, but then calls it with arguments that change
+        in differentiability.
+        """
+        if argnum is None:
+            # Infer which arguments to consider trainable
+            _argnum = _get_argnum(args)
+            # Infer whether to unpack from the infered argnum
+            unpack = len(_argnum) == 1
+        else:
+            # For a single integer as argnum, unpack the Jacobian tuple
+            unpack = isinstance(argnum, int)
+            _argnum = [argnum] if unpack else argnum
+
+        jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
+
+        return jac[0] if unpack else jac
 
     return _jacobian_function