Pytorch style dynamic computation graph

[wcqc]

Hi,

Can TensorCircuit be used to construct hybrid classical quantum computation graphs dynamically like is standard in PyTorch?

Thanks!

[refraction-ray]

I am not sure about your use case, could you explain the feature you desire in more detail?
Could https://tensorcircuit.readthedocs.io/en/latest/tutorials/torch_qml.html or https://github.com/tencent-quantum-lab/tensorcircuit/blob/master/examples/hybrid_gpu_pipeline.py help?

[wcqc]

For the above, all l = 6 inputs have a length of n = 4, how about for l = 6 inputs each of which has a different length? And for each I only want to do a measurement at the end?

[refraction-ray]

how about for l = 6 inputs each of which has a different length?

You mean by input a RaggedTensor? Otherwise how can we combine the different shape as an input without padding?

And for each I only want to do a measurement at the end?

If you describe the exact user case you want to implement, I may better help you. For now, I am still not 100% sure on the user scenario and what I should implement

[wcqc]

Suppose we have 3 sequences:

seq1: abc
seq2: abcdef
seq3: abcd

then we pad (with X):

seq1: abcXXX
seq2: abcdef
seq3: abcdXX

then we feed all the 3 sequences to the same sequential PQC; however, when doing backprop, for seq1, it should backprop from the measurement (say PauliZ) at the letter c, for seq2, from f, and for seq3, from d. This is possible in this lib (https://github.com/mit-han-lab/torchquantum)

[refraction-ray]

ah, now I see your point. It can be done in tc. It can even be jitted as a static computational graph! which mean it runs much faster than pytorch. I'll present the code snippet soon.

[refraction-ray]

@wcqc , here you go, pls check whether the following is what you desire

import numpy as np
import tensorcircuit as tc

K = tc.set_backend("tensorflow")
n = 4

def qmodule(inp, param, pos):
    c = tc.Circuit(n)
    for i in range(n):
        c.rx(i, theta=inp[i])
    for i in range(n-1):
        c.cx(i, i+1)
    for i in range(n):
        c.rx(i, theta=param[i])
    return K.real(tc.templates.measurements.any_measurements(c, pos, onehot=True))

vgf = K.jit(K.vectorized_value_and_grad(qmodule, argnums=1, vectorized_argnums=(0, 2)))

inps = K.convert_to_tensor(np.array([[0.5, 0.5, 0, 0, 0],
                                    [0.3, 0.4, 0.5, 0, 0],
                                    [0.1, -0.2, 0.3, 0.8, 1.2]]))

# batch (3) of sequence input with different length (2, 3, 5)

poss = K.convert_to_tensor(np.array([[0, 3, 0, 0, 0], 
                                   [0, 0, 3, 0, 0],
                                   [0, 0, 0, 0, 3]]))
# measure Z (3 for Z) on the 2, 3, 5 qubit respectively

param = K.ones([n], dtype="float32")

# random parameters to be trained and differentiated

vgf(inps, param, poss)

perfect jittable and thus much faster than pytorch related implementation. Though you can still wrap the whole thing with a pytorch interface so that you can hybrid the code with torch

[wcqc]

@refraction-ray

I understand the pros of JIT (but once a graph is jitted, it becomes static? ), and there are also use cases for dynamic graphs which are handy.

I changed the PyTorch MNIST example in TC's tutorial to do the following: note the new function "def qpred_1step" and the for loop in "def forward" of "class QuantumNet". After these changes, the following code snippet is closer to how I would do a similar PQC in TorchQuantum, but not sure if this is a correct usage in TensorCircuit.

For example, since the circuit c is defined outside of both QuantumNet and qpred_1step, e.g., would each call of qpred_1step still add gates to the same circuit c? Would the changes made break anything in TensorCircuit.Any comments, suggestions welcome.

(Just testing the logic, and the following runs, but not everything is correct.)

Additionally, PyTorch is in the process of adding native JIT and vmap support, is there any plan for adding native PyTorch support in TensorCircuit w/o needing TensorFlow?

import time
import numpy as np
import tensorflow as tf
import torch

import tensorcircuit as tc

K = tc.set_backend("tensorflow")

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train[..., np.newaxis] / 255.0


def filter_pair(x, y, a, b):
    keep = (y == a) | (y == b)
    x, y = x[keep], y[keep]
    y = y == a
    return x, y


x_train, y_train = filter_pair(x_train, y_train, 1, 5)
x_train_small = tf.image.resize(x_train, (3, 3)).numpy()
x_train_bin = np.array(x_train_small > 0.5, dtype=np.float32)
x_train_bin = np.squeeze(x_train_bin).reshape([-1, 9])
y_train_torch = torch.tensor(y_train, dtype=torch.float32)
x_train_torch = torch.tensor(x_train_bin)
x_train_torch.shape, y_train_torch.shape

n = 9
nlayers = 3

# We define the quantum function,
# note how this function is running on tensorflow

c = tc.Circuit(9)

def qpred_1step(x, weights): # each call of this constructs one layer of c
    for k in range(n):
        c.rx(k, theta=x[k])
    for k in range(n - 1):
        c.cnot(k, k + 1)
    for k in range(n):
        c.rx(k, theta=weights[2 * 0, k])
        c.ry(k, theta=weights[2 * 0 + 1, k])
    ypred = c.expectation_ps(z=[n // 2])
    ypred = K.real(ypred)
    return K.sigmoid(ypred)


# Wrap the function into pytorch form but with tensorflow speed!
qpred_torch = tc.interfaces.torch_interface(qpred_1step, jit=True)


class QuantumNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.q_weights = torch.nn.Parameter(torch.randn([2 * nlayers, n]))

    def forward(self, inputs):
        for j in range(nlayers): # new for loop here
            ypred = qpred_torch(inputs, self.q_weights)
        return ypred


net = QuantumNet()
# net(x_train_torch[0])

criterion = torch.nn.BCELoss()
opt = torch.optim.Adam(net.parameters(), lr=1e-2)
nepochs = 500
nbatch = 32
times = []

for epoch in range(nepochs):
    index = np.random.randint(low=0, high=100, size=nbatch)
    # index = np.arange(nbatch)
    inputs, labels = x_train_torch[index], y_train_torch[index]
    opt.zero_grad()

    with torch.set_grad_enabled(True):
        time0 = time.time()
        yps = []
        for i in range(nbatch):
            yp = net(inputs[i])
            yps.append(yp)
        yps = torch.stack(yps)
        loss = criterion(
                torch.reshape(yps, [nbatch, 1]), torch.reshape(labels, [nbatch, 1])
        )
        loss.backward()
        if epoch % 100 == 0:
            print(loss)
        opt.step()
        time1 = time.time()
        times.append(time1 - time0)

print("training time per step: ", np.mean(time1 - time0))

[refraction-ray]

and there are also use cases for dynamic graphs which are handy

yes, many of these use cases can be transformed into static graph smartly, and if not, of course you can just use dynamic graph as long as you don't dress function with @K.jit, all backends (tf, jax, torch) by default support such dynamic graph without jit.

Additionally, PyTorch is in the process of adding native JIT and vmap support, is there any plan for adding native PyTorch support in TensorCircuit w/o needing TensorFlow?

One can now also run on PyTorch backend if speed is not the concern (no vmap and jit though). Last time I try jit and vmap feature of pytorch, they are still buggy and not robust to use. If these features are mature, it is actually very easy to integrate into tensorcircuit (less than 100 lines of code).

---

Still not very sure of the code,

for j in range(nlayers): # new for loop here
            ypred = qpred_torch(inputs, self.q_weights)

This one is not a loop right? Because ypred is not sent to anywhere, and the loop for nlayer times make no sense. You mean ypred=qpred_torch(ypred, self.q_weights)? I want to double-check the thing you want - is it to run the same circuit module as an RNN module or each time the circuit is called it gets deeper?

would each call of qpred_1step still add gates to the same circuit c

namely, you want this or want to avoid this?

[refraction-ray]

what do you mean by one coherent circuit, again confused about the exact use case you desire. Could you explain what you want to implement in rough pseudo code? And if c is outside the function, it means that the circuit is adding gates for each call which is not reasonable to me when you do a training with several hundred forward calls (hundred of layers for the circuit then??).

[wcqc]

it means that the circuit is adding gates for each call which is not reasonable to me when you do a training with several hundred forward calls (hundred of layers for the circuit then??).

Yes, this is what is desired, ie adding gates to the SAME circuit with each call of qpred_1step, instead of creating a new circuit each time. Would this work with tensorcircuit?

[wcqc]

@refraction-ray pls see some pseudo code below

[wcqc]

To me, your issue is that the int passed to qpred_torch will fail the program.

One workaround with dynamic graph is just using pure pytorch backend as below, then no torch_interface is required (can be slow but I think at least as fast as torchquantum, since they also use torch):

import tensorcircuit as tc
from functools import partial

K = tc.set_backend("pytorch")
n = 3
def qpred_1step(x, weights, depth): # each call of this constructs one layer of c
    c = tc.Circuit(n)
    for k in range(n):
        c.rx(k, theta=x[k])
    for k in range(n - 1):
        c.cnot(k, k + 1)
    for j in range(depth):
        for k in range(n):
            c.rx(k, theta=weights[2 * j, k])
            c.ry(k, theta=weights[2 * j + 1, k])
    ypred = c.expectation_ps(z=[n // 2])
    ypred = K.real(ypred)
    return K.sigmoid(ypred)

qpred_1step(K.ones([n]), K.ones([4, n]), 2)

remember to update to the github master version of tensorcircuit since there is a sigmoid bug on torch backed I just identified.

Another way to keep the jit on tf backend while allow int in torch function is as follows:

import tensorcircuit as tc
from functools import partial

K = tc.set_backend("tensorflow")
n = 3
def qpred_1step(x, weights, depth): # each call of this constructs one layer of c
    c = tc.Circuit(n)
    for k in range(n):
        c.rx(k, theta=x[k])
    for k in range(n - 1):
        c.cnot(k, k + 1)
    for j in range(depth):
        for k in range(n):
            c.rx(k, theta=weights[2 * j, k])
            c.ry(k, theta=weights[2 * j + 1, k])
    ypred = c.expectation_ps(z=[n // 2])
    ypred = K.real(ypred)
    return K.sigmoid(ypred)

qpred_torch_dict = {}

def qpred_torch(x, weights, depth):
    if depth not in qpred_torch_dict:
        pf = partial(qpred_1step, depth=depth)
        qpred_torch_dict[depth] = tc.interfaces.torch_interface(pf, jit=True)
    return qpred_torch_dict[depth](x, weights)
  
qpred_torch(torch.ones([n]), torch.ones([3*2, n]), 3)

Help that helps.

If pure Pytorch backend is used like your example code snippet above, would it be possible to make qpred_torch above take an input x of size (batch_size, ...) where the 1st dim is the batch_size, and when returning from qpredc_torch a tensor of size (batch_size, ...) is returned? Or do we have to do this via vmap, as in this example: https://tensorcircuit.readthedocs.io/en/latest/tutorials/torch_qml.html#Batched-Version

[refraction-ray]

yes, there is no native vmap support for torch backend, two workarounds if you dont want to use torch interface and other backends that support vmap:

1. wrap the torch function with a naive "for" for the first batch dimension (could be very slow)
2. try wrap the torch function with torch experimental vmap (try on your own, I am not familiar with torch + vmap, but there is such a thing)

[wcqc]

Here is some pseudocode:

import time
import numpy as np
import tensorflow as tf
import torch

import tensorcircuit as tc

K = tc.set_backend("pytorch")

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train[..., np.newaxis] / 255.0


def filter_pair(x, y, a, b):
    keep = (y == a) | (y == b)
    x, y = x[keep], y[keep]
    y = y == a
    return x, y


x_train, y_train = filter_pair(x_train, y_train, 1, 5)
x_train_small = tf.image.resize(x_train, (3, 3)).numpy()
x_train_bin = np.array(x_train_small > 0.5, dtype=np.float32)
x_train_bin = np.squeeze(x_train_bin).reshape([-1, 9])
y_train_torch = torch.tensor(y_train, dtype=torch.float32)
x_train_torch = torch.tensor(x_train_bin)
x_train_torch.shape, y_train_torch.shape

n = 9


# We define the quantum function,
# note how this function is running on tensorflow



def qpred_1step(x, c): # each call of this constructs one layer of c
    for k in range(n):
        c.rx(k, theta=x[k])
    for k in range(n - 1):
        c.cnot(k, k + 1)
    ypred = c.expectation_ps(z=[n // 2])
    ypred = K.real(ypred)
    return ypred

class QuantumNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.depth = 20
        self.c = tc.Circuit(9)

    def forward(self, inputs):
        ypreds = []
        for j in range(self.depth): # new for loop here
            ypreds.append(qpred_1step(inputs, self.c))
        return ypreds


net = QuantumNet()
# net(x_train_torch[0])

criterion = torch.nn.BCELoss()
opt = torch.optim.Adam(net.parameters(), lr=1e-2)
nepochs = 500
nbatch = 32
times = []

for epoch in range(nepochs):
    index = np.random.randint(low=0, high=100, size=nbatch)
    # index = np.arange(nbatch)
    inputs, labels = x_train_torch[index], y_train_torch[index]
    opt.zero_grad()

    with torch.set_grad_enabled(True):
        time0 = time.time()
        yps = []
        for i in range(nbatch):
            yp = net(inputs[i])
            yps.append(yp)
        yps = torch.stack(yps)
        loss = criterion(
                torch.reshape(yps, [nbatch, 1]), torch.reshape(labels, [nbatch, 1])
        )
        loss.backward()
        if epoch % 100 == 0:
            print(loss)
        opt.step()
        time1 = time.time()
        times.append(time1 - time0)

print("training time per step: ", np.mean(time1 - time0))

[refraction-ray]

How would one implement discarding with density matrix simulation

Since tensor network engine is not like a state simulator, my suggestion is to directly use 4 qubits for the circuit in the figure (no extra cost for this), then the state after discarding the qubits is the reduced_density_matrix as you presented. If the final output you want is some Pauli string expectation, then just compute it by c.expectation_ps(), no need to partial trace by reduced_density_matrix

Your following solution is nearly right though I dont recommend this approach. s1_reduced is a 2-qubit density matrix so it must be further outer product with the density matrix on the new qubit as the inputs for DMCircuit. Also since the s1_reduced is a density matrix, you must use c2=tc.DMCircuit. And there is no need to append c2 to c1, just compute the expectation over c2 final state which is the correct result. Overall, I prefer the first approach -- just create as many qubits you need instead of trying to save qubits (unlike convenstional state simulators, no need for tensor network simulator to save qubits)

[refraction-ray]

Yes, this is what is desired, ie adding gates to the SAME circuit with each call of qpred_1step, instead of creating a new circuit each time. Would this work with tensorcircuit?

According to your remarks and the presudocode, what you desire is to attach one layer on the exsisting circuit in each training step? This is somehow weird, and most importantly, in the code npeochs=500, which indicates that after the training you will have a circuit with 500 layers? In this case, since everytime you evaluate a totally different circuit (different depth), there is indeed no good way to jit (static graph approach) though I believe the following remark is still valid. But I do think the simulation efficiency with dynamic graph on 500-layer circuit is not going to be very good.

"yes, many of these use cases can be transformed into static graph smartly, and if not, of course you can just use dynamic graph as long as you don't dress function with @K.jit, all backends (tf, jax, torch) by default support such dynamic graph without jit."

[wcqc]

How would one implement discarding with density matrix simulation

Since tensor network engine is not like a state simulator, my suggestion is to directly use 4 qubits for the circuit in the figure (no extra cost for this), then the state after discarding the qubits is the reduced_density_matrix as you presented. If the final output you want is some Pauli string expectation, then just compute it by c.expectation_ps(), no need to partial trace by reduced_density_matrix

Would this solution become inefficient if discarding is done say 100 times? Since then we need to initialise a > 100 qubit initial state?

Your following solution is nearly right though I dont recommend this approach. s1_reduced is a 2-qubit density matrix so it must be further outer product with the density matrix on the new qubit as the inputs for DMCircuit. Also since the s1_reduced is a density matrix, you must use c2=tc.DMCircuit. And there is no need to append c2 to c1, just compute the expectation over c2 final state which is the correct result. Overall, I prefer the first approach -- just create as many qubits you need instead of trying to save qubits (unlike convenstional state simulators, no need for tensor network simulator to save qubits)

If c2 is not appended to c1, when doing backprop, would the gradients from c2 flow back to and affect c1 or not? in other words if c2 is not appended to c1 do c2 and c1 become disconnected, which is NOT desired in my case, as I want them to be connected, in the sense that the remaining qubits of c1 after discarding are kept by c2, and this continues to c_k, eg for k=100, and all of c1 to c_k need to connected — ie they are on one tensor network.

Will provide code snippet to further clarify.

[wcqc]

Yes, this is what is desired, ie adding gates to the SAME circuit with each call of qpred_1step, instead of creating a new circuit each time. Would this work with tensorcircuit?

According to your remarks and the presudocode, what you desire is to attach one layer on the exsisting circuit in each training step? This is somehow weird, and most importantly, in the code npeochs=500, which indicates that after the training you will have a circuit with 500 layers? In this case, since everytime you evaluate a totally different circuit (different depth), there is indeed no good way to jit (static graph approach) though I believe the following remark is still valid. But I do think the simulation efficiency with dynamic graph on 500-layer circuit is not going to be very good.

"yes, many of these use cases can be transformed into static graph smartly, and if not, of course you can just use dynamic graph as long as you don't dress function with @K.jit, all backends (tf, jax, torch) by default support such dynamic graph without jit."

Will provide code snippet to further clarify.

[refraction-ray]

Would this solution become inefficient if discarding is done say 100 times? Since then we need to initialise a > 100 qubit initial state?

In principle, no negative effect as claim 100 qubits is nearly the same as claiming 1 qubits, no extra time cost. But you can always have some micro benchmark to see whether this is true.

If c2 is not appended to c1, when doing backprop, would the gradients from c2 flow back to and affect c1 or not?

Yes, backprop will work as expected as the intermediate variable s1_reduced is also a node of the computational graph.