Repository: PennyLaneAI/pennylane
Pull request #1911 (2021-11-17 18:44:35)

Lie gradient flow optimizer
**Context:**
Implementation of Lie gradient flow on quantum circuits.

**Description of the Change:**
New optimizer is added

The lie gradient optimizer grows the circuit by appending unitaries corresponding to the Lie gradient to the circuit.
This is currently done via a @qml.qfunc_transform and this works reasonably well. I wrote a couple of tests to make sure that calling opt.step() works and that what is being calculated is checked against what you would get with doing everything in numpy.

One important issue that I'm having is calculating the cost after the circuit has been transformed.

For some reason, line 236 (now commented out):

cost_fn = qml.map(circuit, self.observables, device=self.circuit.device, measure='expval')([],[])

throws the following error
```python
 _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../pennylane/optimize/lie_gradient.py:258: in step
    self.step_and_cost(
../../pennylane/optimize/lie_gradient.py:236: in step_and_cost
    cost_fn = qml.map(circuit, self.observables, device=self.circuit.device, measure='expval')([],[])
../../pennylane/collections/qnode_collection.py:277: in __call__
    results = self.evaluate(args, kwargs)
../../pennylane/collections/qnode_collection.py:233: in evaluate
    results.append(q(*args, **kwargs))
../../pennylane/qnode.py:558: in __call__
    self.construct(args, kwargs)
../../pennylane/qnode.py:494: in construct
    self._qfunc_output = self.func(*args, **kwargs)
../../pennylane/collections/map.py:137: in circuit
    return MEASURE_MAP[_m](_obs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

op = tensor([], dtype=float64, requires_grad=True)

    def expval(op):
        r"""Expectation value of the supplied observable.
    
        **Example:**
    
        .. code-block:: python3
    
            dev = qml.device("default.qubit", wires=2)
    
            @qml.qnode(dev)
            def circuit(x):
                qml.RX(x, wires=0)
                qml.Hadamard(wires=1)
                qml.CNOT(wires=[0, 1])
                return qml.expval(qml.PauliY(0))
    
        Executing this QNode:
    
        >>> circuit(0.5)
        -0.4794255386042029
    
        Args:
            op (Observable): a quantum observable object
    
        Raises:
            QuantumFunctionError: `op` is not an instance of :class:`~.Observable`
        """
        if not isinstance(op, (Observable, qml.Hamiltonian)):
            raise qml.QuantumFunctionError(
>               "{} is not an observable: cannot be used with expval".format(op.name)
            )
E           AttributeError: 'tensor' object has no attribute 'name'

../../pennylane/measure.py:255: AttributeError
```
It has something todo with fact that products of Paulis are considered tensors and not Observables.

I'm sure there is a better way of doing what I'm doing now, so I look forward to your thoughts @josh146 .


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1911](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (2e81e05) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/7e69c099bfa19de40b5f74157e9ff7895ef97cb6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7e69c09) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1911      +/-   ##
==========================================
+ Coverage   98.84%   98.86%   +0.01%     
==========================================
  Files         219      221       +2     
  Lines       16802    17086     +284     
==========================================
+ Hits        16608    16892     +284     
  Misses        194      194              
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/optimize/lie\_gradient.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wdGltaXplL2xpZV9ncmFkaWVudC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/fourier/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvX19pbml0X18ucHk=) | `100.00% <0.00%> (ø)` | |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `100.00% <0.00%> (ø)` | |
| [pennylane/ops/qubit/matrix\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9tYXRyaXhfb3BzLnB5) | `99.18% <0.00%> (ø)` | |
| [pennylane/ops/qubit/parametric\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9wYXJhbWV0cmljX29wcy5weQ==) | `99.83% <0.00%> (ø)` | |
| [pennylane/templates/embeddings/angle.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9lbWJlZGRpbmdzL2FuZ2xlLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/ops/qubit/state\_preparation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9zdGF0ZV9wcmVwYXJhdGlvbi5weQ==) | `100.00% <0.00%> (ø)` | |
| [pennylane/fourier/reconstruct.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvcmVjb25zdHJ1Y3QucHk=) | `100.00% <0.00%> (ø)` | |
| [pennylane/math/single\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvc2luZ2xlX2Rpc3BhdGNoLnB5) | `99.57% <0.00%> (+0.03%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [7e69c09...2e81e05](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1911?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: [sc-9759]
therooler: I implemented the approximate and exact (non-trotterized) lie gradient. The optimizer takes a single arg to turn of exact or Trotterized evolution, and a restriction of the lie gradient to some subspace can now be achieved by passing a qml.Hamitonian object that spans this subspace. I still need to add more tests to ensure that everything is working properly, but after that it should be ready for another review I think.
therooler: Hi, I think I managed to address most of your comments:
- I figured out a way to flatten the tapes and vectorize everything and it works. Yay!
- I improved the docs and rendered it to make sure that things look all right (sorry for the horrible state it was in before :see_no_evil: )
- Added a better description of the optimizer, probably needs improvement.

Thanks for the great feedback :smile: .
therooler: I checked the docs, did the pylint and black. I think it all looks good! Oh and I changed everything to Lie algebra optimizer with Riemannian gradients.

Patch:
diff --git a/doc/introduction/optimizers.rst b/doc/introduction/optimizers.rst
index a1f3594cd..97d4aee2b 100644
--- a/doc/introduction/optimizers.rst
+++ b/doc/introduction/optimizers.rst
@@ -18,8 +18,8 @@ NumPy Interface
 ^^^^^^^^^^^^^^^
 
 When using the standard NumPy interface, PennyLane offers some built-in optimizers.
-Some of these are specific to quantum optimization, such as the :mod:`~.QNGOptimizer`,
-:mod:`~.RotosolveOptimizer`, :mod:`~.RotoselectOptimizer`, and :mod:`~.ShotAdaptiveOptimizer`.
+Some of these are specific to quantum optimization, such as the :class:`~.QNGOptimizer`, :class:`~.LieAlgebraOptimizer`
+:class:`~.RotosolveOptimizer`, :class:`~.RotoselectOptimizer`, and :class:`~.ShotAdaptiveOptimizer`.
 
 :html:`<div class="summary-table">`
 
@@ -29,6 +29,7 @@ Some of these are specific to quantum optimization, such as the :mod:`~.QNGOptim
     ~pennylane.AdagradOptimizer
     ~pennylane.AdamOptimizer
     ~pennylane.GradientDescentOptimizer
+    ~pennylane.LieAlgebraOptimizer
     ~pennylane.MomentumOptimizer
     ~pennylane.NesterovMomentumOptimizer
     ~pennylane.QNGOptimizer
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 96a8b8684..ba7a955e7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* Added `qml.LieAlgebraOptimizer`, a new quantum-aware Lie Algebra optimizer that allows
+  one to perform gradient descent on the special unitary group.
+  [(#1911)](https://github.com/PennyLaneAI/pennylane/pull/1911)
+  
 * Added functions for computing the values of atomic and molecular orbitals at a given position.
   [(#1867)](https://github.com/PennyLaneAI/pennylane/pull/1867)
 
@@ -607,4 +611,4 @@ This release contains contributions from (in alphabetical order):
 Guillermo Alonso-Linaje, Juan Miguel Arrazola, Samuel Banning, Benjamin Cordier, Alain Delgado,
 Olivia Di Matteo, David Ittah, Josh Izaac, Soran Jahangiri, Jalani Kanem, Ankit Khandelwal, Shumpei
 Kobayashi, Robert Lang, Christina Lee, Cedric Lin, Alejandro Montanez, Romain Moyard, Antal Száva,
-Maria Schuld, Jay Soni, Rodrigo Vargas, David Wierichs
+Maria Schuld, Jay Soni, Rodrigo Vargas, David Wierichs, Roeland Wiersema
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index b774020c1..345c7988f 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -48,7 +48,6 @@ from pennylane.templates.embeddings import *
 from pennylane.templates.layers import *
 from pennylane.templates.state_preparations import *
 from pennylane.templates.subroutines import *
-from pennylane.optimize import *
 from pennylane import qaoa
 from pennylane.qnode import QNode, qnode
 import pennylane.qnode_old
@@ -69,6 +68,7 @@ from pennylane.transforms import (
     quantum_monte_carlo,
     apply_controlled_Q,
 )
+from pennylane.optimize import *
 from pennylane.vqe import ExpvalCost, VQECost
 
 # QueuingContext and collections needs to be imported after all other pennylane imports
diff --git a/pennylane/optimize/__init__.py b/pennylane/optimize/__init__.py
index 969f38568..425596a0a 100644
--- a/pennylane/optimize/__init__.py
+++ b/pennylane/optimize/__init__.py
@@ -20,6 +20,7 @@ This module contains optimizers for the standard :mod:`QNode` class, which uses
 from .adagrad import AdagradOptimizer
 from .adam import AdamOptimizer
 from .gradient_descent import GradientDescentOptimizer
+from .lie_algebra import LieAlgebraOptimizer
 from .momentum import MomentumOptimizer
 from .nesterov_momentum import NesterovMomentumOptimizer
 from .rms_prop import RMSPropOptimizer
@@ -34,6 +35,7 @@ __all__ = [
     "AdagradOptimizer",
     "AdamOptimizer",
     "GradientDescentOptimizer",
+    "LieAlgebraOptimizer",
     "MomentumOptimizer",
     "NesterovMomentumOptimizer",
     "RMSPropOptimizer",
diff --git a/pennylane/optimize/lie_algebra.py b/pennylane/optimize/lie_algebra.py
new file mode 100644
index 000000000..9c8a172eb
--- /dev/null
+++ b/pennylane/optimize/lie_algebra.py
@@ -0,0 +1,363 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Lie algebra optimizer"""
+import warnings
+
+import numpy as np
+from scipy.sparse.linalg import expm
+import pennylane as qml
+
+
+@qml.qfunc_transform
+def append_time_evolution(tape, riemannian_gradient, t, n, exact=False):
+    r"""Append an approximate time evolution, corresponding to a Riemannian
+    gradient on the Lie group, to an existing circuit.
+
+    We want to implement the time evolution generated by an operator of the form
+
+    .. math::
+
+        \text{grad} f(U) = sum_i c_i O_i,
+
+    where :math:`O_i` are Pauli words and :math:`c_t \in \mathbb{R}`.
+    If ``exact`` is ``False``, we Trotterize this operator and apply the unitary
+
+    .. math::
+
+        U' = \prod_{n=1}^{N_{Trot.}} \left(\prod_i \exp{-it / N_{Trot.} O_i}\right),
+
+    which is then appended to the current circuit.
+
+    If ``exact`` is ``True``, we calculate the exact time evolution for the Riemannian gradient
+     by way of the matrix exponential.
+
+    .. math:
+
+        U' = \exp{-it \text{grad} f(U)}
+
+    and append this unitary.
+
+    Args:
+        tape (QuantumTape or .QNode): circuit to transform.
+        riemannian_gradient (.Hamiltonian): Hamiltonian object representing the Riemannian gradient.
+        t (float): time evolution parameter.
+        n (int): number of Trotter steps.
+
+    """
+    for obj in tape.operations:
+        qml.apply(obj)
+    if exact:
+        qml.QubitUnitary(
+            expm(-1j * t * qml.utils.sparse_hamiltonian(riemannian_gradient).toarray()),
+            wires=range(len(riemannian_gradient.wires)),
+        )
+    else:
+        qml.templates.ApproxTimeEvolution(riemannian_gradient, t, n)
+
+    for obj in tape.measurements:
+        qml.apply(obj)
+
+
+@qml.batch_transform
+def algebra_commutator(tape, observables, lie_algebra_basis_names, nqubits):
+    """Calculate the Riemannian gradient in the Lie algebra with the parameter shift rule
+    (see :meth:`LieAlgebraOptimizer.get_omegas`).
+
+    Args:
+        tape (.QuantumTape or .QNode): input circuit
+        observables (list[.Observable]): list of observables to be measured. Can be grouped.
+        lie_algebra_basis_names (list[str]): List of strings corresponding to valid Pauli words.
+        nqubits (int): the number of qubits.
+
+    Returns:
+         func: Function which accepts the same arguments as the QNode. When called, this
+         function will return the Lie algebra commutator.
+
+    """
+    tapes_plus_total = []
+    tapes_min_total = []
+    for obs in observables:
+        for o in obs:
+            # create a list of tapes for the plus and minus shifted circuits
+            tapes_plus = [qml.tape.JacobianTape(p + "_p") for p in lie_algebra_basis_names]
+            tapes_min = [qml.tape.JacobianTape(p + "_m") for p in lie_algebra_basis_names]
+
+            # loop through all operations on the input tape
+            for op in tape.operations:
+                for t in tapes_plus + tapes_min:
+                    with t:
+                        qml.apply(op)
+            for i, t in enumerate(tapes_plus):
+                with t:
+                    qml.PauliRot(
+                        np.pi / 2,
+                        lie_algebra_basis_names[i],
+                        wires=list(range(nqubits)),
+                    )
+                    qml.expval(o)
+            for i, t in enumerate(tapes_min):
+                with t:
+                    qml.PauliRot(
+                        -np.pi / 2,
+                        lie_algebra_basis_names[i],
+                        wires=list(range(nqubits)),
+                    )
+                    qml.expval(o)
+            tapes_plus_total.extend(tapes_plus)
+            tapes_min_total.extend(tapes_min)
+    return tapes_plus_total + tapes_min_total, None
+
+
+class LieAlgebraOptimizer:
+    r"""Lie algebra optimizer.
+
+    Riemannian gradient descent algorithms can be used to optimize a function directly on a Lie group
+    as opposed to on a Euclidean parameter space. Consider the function
+    :math:`f(U) = \text{Tr}(U \rho_0 U^\dagger H)`
+    for a given Hamiltonian :math:`H`, unitary :math:`U\in \text{SU}(2^N)` and initial state
+    :math:`\rho_0`. One can show that this function is minimized by the flow equation
+
+    .. math::
+
+        \dot{U} = \text{grad}f(U)
+
+    where :math:`\text{grad}` is the Riemannian gradient operator on :math:`\text{SU}(2^N)`.
+    By discretizing the flow above, we see that a step of the Lie algebra optimizer
+    iterates the Riemannian gradient flow on :math:`\text{SU}(2^N)` as
+
+    .. math::
+
+        U^{(t+1)} = \exp\left\{\epsilon\: \text{grad}f(U^{(t)}) U^{(t)}\right\},
+
+
+    where :math:`\epsilon` is a user-defined hyperparameter corresponding to the step size.
+
+    The gradient in the Lie algebra is given by
+
+    .. math::
+
+        \text{grad}f(U^{(t)}) = -\left[U \rho U^\dagger, H\right] .
+
+
+    Hence we see that subsequent steps of this optimizer will append the unitary generated by the Lie
+    gradient and grow the circuit.
+
+    The exact Riemannian gradient flow on :math:`\text{SU}(2^N)` has desirable optimization properties
+    that can guarantee convergence to global minima under mild assumptions. However, this comes
+    at a cost. Since :math:`\text{dim}(\text{SU}(2^N)) = 4^N-1`, we need an exponential number
+    of parameters to calculate the gradient. This will not be problematic for small systems (:math:`N<5`),
+    but will quickly get out of control as the number of qubits increases.
+
+    To resolve this issue, we can restrict the Riemannian gradient to a subspace of the Lie algebra and calculate an
+    approximate Riemannian gradient flow. The choice of restriction will affect the optimization behavior
+    and quality of the final solution.
+
+    For more information on Riemannian gradient flows on Lie groups see
+    `T. Schulte-Herbrueggen et. al. (2008) <https://arxiv.org/abs/0802.4195>`_
+    and the application to quantum circuits
+    `Wiersema and Killoran (2021) <https://arxiv.org/abs/>`_.
+
+    Args:
+        circuit (.QNode): a user defined circuit that does not take any arguments and returns
+            the expectation value of a ``qml.Hamiltonian``.
+        stepsize (float): the user-defined hyperparameter :math:`\epsilon`.
+        restriction (.Hamiltonian): Restrict the Lie algebra to a corresponding subspace of
+            the full Lie algebra. This restriction should be passed in the form of a
+            ``qml.Hamiltonian`` that consists only of Pauli words.
+        exact (bool): Flag that indicates wether we approximate the Riemannian gradient with a
+            Trotterization or calculate the exact evolution via a matrix exponential. The latter is
+            not hardware friendly and can only be done in simulation.
+
+    **Examples**
+
+    Define a Hamiltonian cost function to minimize:
+
+    >>> coeffs = [-1., -1., -1.]
+    >>> observables = [qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0) @ qml.PauliX(1)]
+    >>> hamiltonian = qml.Hamiltonian(coeffs, observables)
+
+    Create an initial state and return the expectation value of the Hamiltonian:
+
+    >>> @qml.qnode(qml.device("default.qubit", wires=2))
+    ... def quant_fun():
+    ...     qml.RX(0.1, wires=[0])
+    ...     qml.RY(0.5, wires=[1])
+    ...     qml.CNOT(wires=[0,1])
+    ...     qml.RY(0.6, wires=[0])
+    ...     return qml.expval(hamiltonian)
+
+    Instantiate the optimizer with the initial circuit and the cost function and set the stepsize
+    accordingly:
+
+    >>> opt = qml.LieAlgebraOptimizer(circuit=quant_fun, stepsize=0.1)
+
+    Applying 5 steps gets us close the ground state of :math:`E\approx-2.23`:
+
+    >>> for step in range(6):
+    ...    circuit, cost = opt.step_and_cost()
+    ...    print(f"Step {step} - cost {cost}")
+    Step 0 - cost -1.3351865007304005
+    Step 1 - cost -1.9937887238935206
+    Step 2 - cost -2.1524234485729834
+    Step 3 - cost -2.1955105378898487
+    Step 4 - cost -2.2137628169764256
+    Step 5 - cost -2.2234364822091575
+
+    The optimized circuit is returned at each step, and can be used as any other QNode:
+
+    >>> circuit()
+    -2.2283086057521713
+
+    """
+    # pylint: disable=too-many-arguments
+    # pylint: disable=too-many-instance-attributes
+    def __init__(self, circuit, stepsize=0.01, restriction=None, exact=False, trottersteps=1):
+
+        if not isinstance(circuit, qml.QNode):
+            raise TypeError(f"circuit must be a QNode, received {type(circuit)}")
+
+        self.circuit = circuit
+        self.circuit.construct([], {})
+        self.hamiltonian = circuit.func().obs
+        if not isinstance(self.hamiltonian, qml.Hamiltonian):
+            raise TypeError(
+                f"circuit must return the expectation value of a Hamiltonian,"
+                f"received {type(circuit.func().obs)}"
+            )
+        self.nqubits = len(circuit.device.wires)
+
+        if self.nqubits > 4:
+            warnings.warn(
+                "The exact Riemannian gradient is exponentially expensive in the number of qubits, "
+                f"optimizing a {self.nqubits} qubit circuit may be slow.",
+                UserWarning,
+            )
+        if restriction is not None and not isinstance(restriction, qml.Hamiltonian):
+            raise TypeError(f"restriction must be a Hamiltonian, received {type(restriction)}")
+        (
+            self.lie_algebra_basis_ops,
+            self.lie_algebra_basis_names,
+        ) = self.get_su_n_operators(restriction)
+        self.exact = exact
+        self.trottersteps = trottersteps
+        self.coeffs, self.observables = self.hamiltonian.terms
+        self.stepsize = stepsize
+
+    def step(self):
+        r"""Update the circuit with one step of the optimizer.
+
+        Returns:
+           float: the optimized circuit and the objective function output prior
+           to the step.
+        """
+        return self.step_and_cost()[0]
+
+    def step_and_cost(self):
+        r"""Update the circuit with one step of the optimizer and return the corresponding
+        objective function value prior to the step.
+
+        Returns:
+            tuple[.QNode, float]: the optimized circuit and the objective function output prior
+            to the step.
+        """
+        # pylint: disable=not-callable
+
+        cost = self.circuit()
+        omegas = self.get_omegas()
+        non_zero_omegas = -omegas[omegas != 0]
+
+        nonzero_idx = np.nonzero(omegas)[0]
+        non_zero_lie_algebra_elements = [self.lie_algebra_basis_names[i] for i in nonzero_idx]
+
+        lie_gradient = qml.Hamiltonian(
+            non_zero_omegas,
+            [qml.grouping.string_to_pauli_word(ps) for ps in non_zero_lie_algebra_elements],
+        )
+        new_circuit = append_time_evolution(
+            lie_gradient, self.stepsize, self.trottersteps, self.exact
+        )(self.circuit.func)
+
+        self.circuit = qml.QNode(new_circuit, self.circuit.device)
+        return self.circuit, cost
+
+    def get_su_n_operators(self, restriction):
+        r"""Get the SU(N) operators. The dimension of the group is :math:`N^2-1`.
+
+        Args:
+            restriction (.Hamiltonian): Restrict the Riemannian gradient to a subspace.
+
+        Returns:
+            tuple[list[array[complex]], list[str]]: list of :math:`N^2 \times N^2` NumPy complex arrays and corresponding Pauli words.
+        """
+
+        operators = []
+        names = []
+        # construct the corresponding pennylane observables
+        wire_map = dict(zip(range(self.nqubits), range(self.nqubits)))
+        if restriction is None:
+            for ps in qml.grouping.pauli_group(self.nqubits):
+                operators.append(ps)
+                names.append(qml.grouping.pauli_word_to_string(ps, wire_map=wire_map))
+        else:
+            for ps in set(restriction.ops):
+                operators.append(ps)
+                names.append(qml.grouping.pauli_word_to_string(ps, wire_map=wire_map))
+
+        return operators, names
+
+    def get_omegas(self):
+        r"""Measure the coefficients of the Riemannian gradient with respect to a Pauli word basis.
+        We want to calculate the components of the Riemannian gradient in the Lie algebra
+        with respect to a Pauli word basis. For a Hamiltonian of the form :math:`H = \sum_i c_i O_i`,
+        where :math:`c_i\in\mathbb{R}`, this can be achieved by calculating
+
+        .. math::
+
+            \omega_{i,j} = \text{Tr}(c_i[\rho, O_i] P_j)
+
+        where :math:`P_j` is a Pauli word in the set of Pauli monomials on :math:`N` qubits.
+
+        Via the parameter shift rule, the commutator can be calculated as
+
+        .. math::
+
+            [\rho, O_i] = \frac{1}{2}(V(\pi/2) \rho V^\dagger(\pi/2) - V(-\pi/2) \rho V^\dagger(-\pi/2))
+
+        where :math:`V` is the unitary generated by the Pauli word :math:`V(\theta) = \exp\{-i\theta P_j\}`.
+
+        Returns:
+            array: array of omegas for each direction in the Lie algebra.
+
+        """
+
+        obs_groupings, _ = qml.grouping.group_observables(self.observables, self.coeffs)
+        # get all circuits we need to calculate the coefficients
+        circuits = algebra_commutator(
+            self.circuit.qtape,
+            obs_groupings,
+            self.lie_algebra_basis_names,
+            self.nqubits,
+        )[0]
+        circuits = qml.execute(circuits, self.circuit.device, gradient_fn=None)
+        circuits_plus = np.array(circuits[: len(circuits) // 2]).reshape(
+            len(self.coeffs), len(self.lie_algebra_basis_names)
+        )
+        circuits_min = np.array(circuits[len(circuits) // 2 :]).reshape(
+            len(self.coeffs), len(self.lie_algebra_basis_names)
+        )
+
+        # For each observable O_i in the Hamiltonian, we have to calculate all Lie coefficients
+        omegas = 0.5 * (circuits_plus - circuits_min)
+
+        return np.dot(self.coeffs, omegas)


Test patch:
diff --git a/tests/optimize/test_lie_algebra.py b/tests/optimize/test_lie_algebra.py
new file mode 100644
index 000000000..389185dbd
--- /dev/null
+++ b/tests/optimize/test_lie_algebra.py
@@ -0,0 +1,375 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the ``LieAlgebraOptimizer``.
+"""
+import pytest
+
+from scipy.sparse.linalg import expm
+import numpy as np
+
+import pennylane as qml
+from pennylane.optimize import LieAlgebraOptimizer
+
+
+def circuit_1():
+    """Simple circuit."""
+    qml.Hadamard(wires=[0])
+    qml.Hadamard(wires=[1])
+
+
+def circuit_2():
+    """Simply parameterized circuit."""
+    qml.RX(0.1, wires=[0])
+    qml.RY(0.5, wires=[1])
+    qml.CNOT(wires=[0, 1])
+    qml.RY(0.6, wires=[0])
+
+
+def circuit_3():
+    """Three-qubit circuit."""
+    qml.RY(0.5, wires=[0])
+    qml.RY(0.6, wires=[1])
+    qml.RY(0.7, wires=[2])
+    qml.CNOT(wires=[0, 1])
+    qml.CNOT(wires=[1, 2])
+    qml.RX(-0.6, wires=[0])
+    qml.RX(-0.3, wires=[1])
+    qml.RX(-0.2, wires=[2])
+
+
+hamiltonian_1 = qml.Hamiltonian(
+    coeffs=[-1.0] * 3,
+    observables=[qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0) @ qml.PauliX(1)],
+)
+
+hamiltonian_2 = qml.Hamiltonian(
+    coeffs=[-0.2, 0.3, -0.15],
+    observables=[
+        qml.PauliY(1),
+        qml.PauliZ(0) @ qml.PauliZ(1),
+        qml.PauliX(0) @ qml.PauliX(1),
+    ],
+)
+
+hamiltonian_3 = qml.Hamiltonian(
+    coeffs=[-2.0], observables=[qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliY(2)]
+)
+
+
+@pytest.mark.parametrize(
+    "circuit,hamiltonian",
+    [
+        (circuit_1, hamiltonian_1),
+        (circuit_1, hamiltonian_2),
+        (circuit_2, hamiltonian_1),
+        (circuit_2, hamiltonian_2),
+        (circuit_3, hamiltonian_3),
+    ],
+)
+def test_lie_algebra_omegas(circuit, hamiltonian):
+    """Test that we calculate the Riemannian gradient coefficients Tr{[rho, H] P_j} correctly."""
+    # pylint: disable=no-member
+
+    nqubits = max([max(ps.wires) for ps in hamiltonian.ops]) + 1
+    wires = range(nqubits)
+    dev = qml.device("default.qubit", wires=nqubits)
+
+    @qml.qnode(dev)
+    def get_state():
+        circuit()
+        return qml.state()
+
+    @qml.qnode(dev)
+    def lie_circuit():
+        circuit()
+        return qml.expval(hamiltonian)
+
+    phi = get_state()
+    rho = np.outer(phi, phi.conj())
+    hamiltonian_np = qml.utils.sparse_hamiltonian(hamiltonian, wires).toarray()
+    lie_algebra_np = hamiltonian_np @ rho - rho @ hamiltonian_np
+    opt = LieAlgebraOptimizer(circuit=lie_circuit)
+    ops = opt.get_su_n_operators(None)[0]
+    omegas_np = []
+    for op in ops:
+        op = qml.utils.expand(op.matrix, op.wires, wires)
+        omegas_np.append(-np.trace(lie_algebra_np @ op).imag / 2)
+    omegas = opt.get_omegas()
+    assert np.allclose(omegas, omegas_np)
+
+
+@pytest.mark.parametrize(
+    "circuit,hamiltonian",
+    [
+        (circuit_1, hamiltonian_1),
+        (circuit_1, hamiltonian_2),
+        (circuit_2, hamiltonian_1),
+        (circuit_2, hamiltonian_2),
+        (circuit_3, hamiltonian_3),
+    ],
+)
+def test_lie_algebra_omegas_restricted(circuit, hamiltonian):
+    """Test that we calculate the (restricted) Riemannian gradient coefficients correctly."""
+    # pylint: disable=no-member
+    nqubits = max([max(ps.wires) for ps in hamiltonian.ops]) + 1
+    wires = range(nqubits)
+    dev = qml.device("default.qubit", wires=nqubits)
+
+    @qml.qnode(dev)
+    def get_state():
+        circuit()
+        return qml.state()
+
+    @qml.qnode(dev)
+    def lie_circuit():
+        circuit()
+        return qml.expval(hamiltonian)
+
+    phi = get_state()
+    rho = np.outer(phi, phi.conj())
+    hamiltonian_np = qml.utils.sparse_hamiltonian(hamiltonian, wires).toarray()
+    lie_algebra_np = hamiltonian_np @ rho - rho @ hamiltonian_np
+
+    restriction = qml.Hamiltonian(
+        coeffs=[1.0] * 3,
+        observables=[qml.PauliX(0), qml.PauliY(1), qml.PauliY(0) @ qml.PauliY(1)],
+    )
+
+    opt = LieAlgebraOptimizer(circuit=lie_circuit, restriction=restriction)
+    ops = opt.get_su_n_operators(restriction)[0]
+    omegas_np = []
+    for op in ops:
+        op = qml.utils.expand(op.matrix, op.wires, wires)
+        omegas_np.append(-np.trace(lie_algebra_np @ op).imag / 2)
+    omegas = opt.get_omegas()
+
+    assert np.allclose(omegas, omegas_np)
+
+
+@pytest.mark.parametrize(
+    "circuit,hamiltonian",
+    [
+        (circuit_1, hamiltonian_1),
+        (circuit_1, hamiltonian_2),
+        (circuit_2, hamiltonian_1),
+        (circuit_2, hamiltonian_2),
+    ],
+)
+def test_lie_algebra_evolution(circuit, hamiltonian):
+    """Test that the optimizer produces the correct unitary to append."""
+    # pylint: disable=no-member
+    nqubits = max([max(ps.wires) for ps in hamiltonian.ops]) + 1
+    wires = range(nqubits)
+    dev = qml.device("default.qubit", wires=nqubits)
+
+    @qml.qnode(dev)
+    def get_state():
+        circuit()
+        return qml.state()
+
+    @qml.qnode(dev)
+    def lie_circuit():
+        circuit()
+        return qml.expval(hamiltonian)
+
+    phi = get_state()
+    rho = np.outer(phi, phi.conj())
+    hamiltonian_np = qml.utils.sparse_hamiltonian(hamiltonian, wires).toarray()
+    lie_algebra_np = hamiltonian_np @ rho - rho @ hamiltonian_np
+
+    phi_exact = expm(-0.001 * lie_algebra_np) @ phi
+    rho_exact = np.outer(phi_exact, phi_exact.conj())
+    opt = LieAlgebraOptimizer(circuit=lie_circuit, stepsize=0.001, exact=True)
+    opt.step()
+
+    cost_pl = opt.circuit()
+    cost_exact = np.trace(rho_exact @ hamiltonian_np)
+    assert np.allclose(cost_pl, cost_exact, atol=1e-2)
+
+
+@pytest.mark.parametrize(
+    "circuit,hamiltonian",
+    [
+        (circuit_1, hamiltonian_1),
+        (circuit_1, hamiltonian_2),
+        (circuit_2, hamiltonian_1),
+        (circuit_2, hamiltonian_2),
+        (circuit_3, hamiltonian_3),
+    ],
+)
+def test_lie_algebra_step(circuit, hamiltonian):
+    """Test that we can take subsequent steps with the optimizer."""
+    nqubits = max([max(ps.wires) for ps in hamiltonian.ops]) + 1
+
+    dev = qml.device("default.qubit", wires=nqubits)
+
+    @qml.qnode(dev)
+    def lie_circuit():
+        circuit()
+        return qml.expval(hamiltonian)
+
+    opt = LieAlgebraOptimizer(circuit=lie_circuit)
+    opt.step()
+    opt.step()
+
+
+@pytest.mark.parametrize(
+    "circuit,hamiltonian",
+    [
+        (circuit_1, hamiltonian_1),
+        (circuit_1, hamiltonian_2),
+        (circuit_2, hamiltonian_1),
+        (circuit_2, hamiltonian_2),
+        (circuit_3, hamiltonian_3),
+    ],
+)
+def test_lie_algebra_step_trotterstep(circuit, hamiltonian):
+    """Test that we can take subsequent steps with the optimizer."""
+    nqubits = max([max(ps.wires) for ps in hamiltonian.ops]) + 1
+
+    dev = qml.device("default.qubit", wires=nqubits)
+
+    @qml.qnode(dev)
+    def lie_circuit():
+        circuit()
+        return qml.expval(hamiltonian)
+
+    opt = LieAlgebraOptimizer(circuit=lie_circuit, trottersteps=3)
+    opt.step()
+    opt.step()
+
+
+def test_lie_algebra_circuit_input_1_check():
+    """Test that a type error is raise for non-QNode circuits."""
+
+    def circuit():
+        qml.RY(0.5, wires=0)
+
+    with pytest.raises(TypeError, match="circuit must be a QNode"):
+        LieAlgebraOptimizer(circuit=circuit, stepsize=0.001)
+
+
+def test_lie_algebra_hamiltonian_input_1_check():
+    """Test that a type error is raise for non-QNode circuits."""
+
+    @qml.qnode(qml.device("default.qubit", wires=3))
+    def circuit():
+        qml.RY(0.5, wires=0)
+        return qml.state()
+
+    with pytest.raises(
+        TypeError,
+        match="circuit must return the expectation value of a Hamiltonian",
+    ):
+        LieAlgebraOptimizer(circuit=circuit, stepsize=0.001)
+
+
+def test_lie_algebra_nqubits_check():
+    """Test that we warn if the system is too big."""
+
+    @qml.qnode(qml.device("default.qubit", wires=5))
+    def circuit():
+        qml.RY(0.5, wires=0)
+        return qml.expval(qml.Hamiltonian(coeffs=[-1.0], observables=[qml.PauliX(0)]))
+
+    with pytest.warns(UserWarning, match="The exact Riemannian gradient is exponentially"):
+        LieAlgebraOptimizer(circuit=circuit, stepsize=0.001)
+
+
+def test_lie_algebra_restriction_check():
+    """Test that a type error is raise for non-QNode circuits."""
+
+    @qml.qnode(qml.device("default.qubit", wires=3))
+    def circuit():
+        qml.RY(0.5, wires=0)
+        return qml.expval(qml.Hamiltonian(coeffs=[-1.0], observables=[qml.PauliX(0)]))
+
+    restriction = "not_a_hamiltonian"
+    with pytest.raises(
+        TypeError,
+        match="restriction must be a Hamiltonian",
+    ):
+        LieAlgebraOptimizer(circuit=circuit, restriction=restriction, stepsize=0.001)
+
+
+def test_docstring_example():
+    """Test the docstring example with Trotterized evolution."""
+    hamiltonian = qml.Hamiltonian(
+        coeffs=[-1.0] * 3,
+        observables=[qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0) @ qml.PauliX(1)],
+    )
+
+    @qml.qnode(qml.device("default.qubit", wires=2))
+    def quant_fun():
+        qml.RX(0.1, wires=[0])
+        qml.RY(0.5, wires=[1])
+        qml.CNOT(wires=[0, 1])
+        qml.RY(0.6, wires=[0])
+        return qml.expval(hamiltonian)
+
+    opt = LieAlgebraOptimizer(circuit=quant_fun, stepsize=0.1)
+
+    for step in range(6):
+        circuit, cost = opt.step_and_cost()
+        print(f"Step {step} - cost {cost}")
+    print(circuit())
+
+    assert np.isclose(cost, -2.23, atol=1e-2)
+
+
+def test_docstring_example_exact():
+    """Test that the optimizer works with matrix exponential."""
+
+    hamiltonian = qml.Hamiltonian(
+        coeffs=[-1.0] * 3,
+        observables=[qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0) @ qml.PauliX(1)],
+    )
+
+    @qml.qnode(qml.device("default.qubit", wires=2))
+    def quant_fun():
+        qml.RX(0.1, wires=[0])
+        qml.RY(0.5, wires=[1])
+        qml.CNOT(wires=[0, 1])
+        qml.RY(0.6, wires=[0])
+        return qml.expval(hamiltonian)
+
+    opt = LieAlgebraOptimizer(circuit=quant_fun, stepsize=0.1, exact=True)
+
+    for step in range(6):
+        _, cost = opt.step_and_cost()
+        print(f"Step {step} - cost {cost}")
+    assert np.isclose(cost, -2.23, atol=1e-2)
+
+
+def test_example_shots():
+    """Test that the optimizer works with finite shots."""
+    hamiltonian = qml.Hamiltonian(
+        coeffs=[-1.0] * 3,
+        observables=[qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0) @ qml.PauliX(1)],
+    )
+
+    @qml.qnode(qml.device("default.qubit", wires=2, shots=1000))
+    def quant_fun():
+        qml.RX(0.1, wires=[0])
+        qml.RY(0.5, wires=[1])
+        qml.CNOT(wires=[0, 1])
+        qml.RY(0.6, wires=[0])
+        return qml.expval(hamiltonian)
+
+    opt = LieAlgebraOptimizer(circuit=quant_fun, stepsize=0.1, exact=False)
+
+    for step in range(3):
+        _, cost = opt.step_and_cost()
+        print(f"Step {step} - cost {cost}")