Statevector vjp and support Hamiltonian in adjoint_jacobian

[chaeyeunpark]

Before submitting

Please complete the following checklist when submitting a PR:

• All new features must include a unit test.
  If you've fixed a bug or added code that should be tested, add a test to the
  tests directory!

• All new functions and code must be clearly commented and documented.
  If you do make documentation changes, make sure that the docs build and
  render correctly by running make docs.

• Ensure that the test suite passes, by running make test.

• Add a new entry to the .github/CHANGELOG.md file, summarizing the
  change, and including a link back to the PR.

• Ensure that code is properly formatted by running make format.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

---

Context:

Description of the Change: We now have pennylane_lightning.lightning_ops.adj_diff module that handles all observables related to adjoint differentiation method. We also have adjoint-like computation of the statevecor vjp.

Benefits:

Possible Drawbacks:

Related GitHub Issues: Closes #216

[mlxd]

This looks amazing, great job @chaeyeunpark ! I have nothing else major to add, but a few comments.

The only thing I'd like to see before we get this merged in are some runtime plots comparing the before and after (v0.23.0, v0.24.0-dev, and this).

[mlxd]

Well that is depressing...

[chaeyeunpark]

Yes. Moreover, as we use clang-cl (for Kokkos), we cannot use pmr even for Windows though MSVC supports it.

[mlxd]

Do they have it even in experimental support? We could wrap std::experimental::pmr if so.

[chaeyeunpark]

Yes, that's true. Let me add this to my task list.

[chaeyeunpark]

I benchmark the performance of adjoint diff with expval(ham). This is the code I used:

import pennylane as qml
import pennylane.numpy as np
from numpy.random import default_rng
import timeit
import sys

rng = default_rng(1337)
num_qubits = 10

def create_random_pauli_prod(num_qubits):
    paulis = [qml.Identity, qml.PauliX, qml.PauliY, qml.PauliZ]
    pindices = rng.integers(0, 3, num_qubits)
    terms = [paulis[pindices[i]](i) for i in range(num_qubits)]
    return qml.operation.Tensor(*terms)

def create_hamiltonian(num_qubits, n_terms):
    ham_coeffs = rng.standard_normal(n_terms)
    ham_terms = [create_random_pauli_prod(num_qubits) for i in range(n_terms)]
    return qml.Hamiltonian(ham_coeffs, ham_terms)

dev = qml.device('lightning.qubit', wires=num_qubits)

def bench_func(n_terms):
    ham = create_hamiltonian(num_qubits, n_terms)
    with qml.tape.QuantumTape() as tape:
        qml.RX(0.23, wires=0)
        qml.CNOT(wires=[0,1])
        qml.RY(0.23, wires=1)
        qml.CNOT(wires=[1,2])
        qml.RZ(0.23, wires=2)
        qml.expval(ham)

    w = tape.get_parameters(tranable_only = True)

    def timing_this():
        weights = rng.random(3)
        weights = np.array(weights)
        w[:3] = weights
        tape.set_parameters(w, trainable_only=True)
        return dev.adjoint_jacobian(tape, np.ones(2**num_qubits, requires_grad=False, dtype=np.complex128))

    return timeit.timeit(timing_this, number=100)

for n_terms in [2**k for k in range(2, 9)]:
    print("{}\t{}".format(n_terms, bench_func(n_terms)))

Using this simple benchmark code, I got

master version
4	19.39245017792564
8	19.997620621928945
16	22.724174362956546
32	26.514014064101502
64	33.93342457001563
128	54.66594595997594
256	104.8523806289304

this PR
4	0.30219985102303326
8	0.20092515798751265
16	0.445317056030035
32	1.591180496965535
64	8.25559274898842
128	19.637774570961483
256	34.774507706053555

The first column is the number of terms in the Hamiltonian.

[chaeyeunpark]

I also get slightly faster results for a simple observables:

import pennylane as qml
import pennylane.numpy as np
from numpy.random import default_rng
import timeit
import sys

rng = default_rng(1337)

def bench_func(num_qubits):
    dev = qml.device('lightning.qubit', wires=num_qubits)
    with qml.tape.QuantumTape() as tape:
        qml.RX(0.23, wires=0)
        qml.CNOT(wires=[0,1])
        qml.RY(0.23, wires=1)
        qml.CNOT(wires=[1,2])
        qml.RZ(0.23, wires=2)
        qml.expval(qml.PauliX(0) @ qml.PauliY(1))

    def timing_this():
        weights = rng.random(3)
        weights = np.array(weights)
        tape.set_parameters(weights, trainable_only=True)
        return dev.adjoint_jacobian(tape, np.ones(2**num_qubits, requires_grad=False, dtype=np.complex128))

    return timeit.timeit(timing_this, number=100)

for num_qubits in [4, 6, 8, 10, 12, 14, 16, 18]:
    print("{}\t{}".format(num_qubits, bench_func(num_qubits)))

And the result is:

master:
4	0.4656622860347852
6	0.2100854809395969
8	0.05818713700864464
10	0.6123739270260558
12	0.2523528659949079
14	0.19010203890502453
16	1.332373121054843
18	4.672340785968117

this PR:
4	0.12154659803491086
6	0.06435875396709889
8	0.061190371052362025
10	0.05885270098224282
12	0.17519149102736264
14	0.20414239005185664
16	0.935425164992921
18	4.56479615601711

The first column is the number of qubits.

[mlxd]

I also get slightly faster results for a simple observables:

import pennylane as qml
import pennylane.numpy as np
from numpy.random import default_rng
import timeit
import sys

rng = default_rng(1337)

def bench_func(num_qubits):
    dev = qml.device('lightning.qubit', wires=num_qubits)
    with qml.tape.QuantumTape() as tape:
        qml.RX(0.23, wires=0)
        qml.CNOT(wires=[0,1])
        qml.RY(0.23, wires=1)
        qml.CNOT(wires=[1,2])
        qml.RZ(0.23, wires=2)
        qml.expval(qml.PauliX(0) @ qml.PauliY(1))

    def timing_this():
        weights = rng.random(3)
        weights = np.array(weights)
        tape.set_parameters(weights, trainable_only=True)
        return dev.adjoint_jacobian(tape, np.ones(2**num_qubits, requires_grad=False, dtype=np.complex128))

    return timeit.timeit(timing_this, number=100)

for num_qubits in [4, 6, 8, 10, 12, 14, 16, 18]:
    print("{}\t{}".format(num_qubits, bench_func(num_qubits)))

And the result is:

master:
4	0.4656622860347852
6	0.2100854809395969
8	0.05818713700864464
10	0.6123739270260558
12	0.2523528659949079
14	0.19010203890502453
16	1.332373121054843
18	4.672340785968117

this PR:
4	0.12154659803491086
6	0.06435875396709889
8	0.061190371052362025
10	0.05885270098224282
12	0.17519149102736264
14	0.20414239005185664
16	0.935425164992921
18	4.56479615601711

The first column is the number of qubits.

Those numbers are amazing --- excellent work putting this together!

[mlxd]

Nothing more to add. Great work, and thanks @chaeyeunpark

[AmintorDusko]

Suggested change

	custom_decomps={'DoubleExcitation': qml.DoubleExcitation.compute_decomposition}
	dev_lig = qml.device("lightning.qubit", wires=qubits, c_dtype = c_dtype,
	custom_decomps=custom_decomps)
	dev_lig = qml.device("lightning.qubit", wires=qubits, c_dtype = c_dtype)

[AmintorDusko]

This custom decomposition is not necessary anymore.

[AmintorDusko]

Excellent work! You have my 100% OK after the small fix, I suggest.

[chaeyeunpark]

Thanks @AmintorDusko and @mlxd for the reviews. I just added one small test. I will merge this PR when Tidy is satisfied with the code.