Autograd tensor as class property leads to issues with automatic differentiation

I’m constructing a VQC class using Pennylane. All’s working well save one exception – I’m finding that having the tensor of parameters as a property of the class is causing issues with automatic differentiation.

Main Question: Do tensors in autograd’s computational graph need to be in globals()? If not, how can i explicitly pass the parameter tensor to the optimization function s.t. autograd is able to use it? Below is a stripped down version of my code showing how the references to the parameters (self.theta), circuit, and qnode are being passed around:

class QuantumClassifier(object):

    def __init__(self, model_config:dict={}):
        self.n_layers = model_config["n_layers"]
        self.m_qubits = model_config["m_qubits"]
        # ...other properties defined here
        self.theta_shape = (self.n_layers, self.m_qubits, 3)
        self.theta = (np.random.randn(*self.theta_shape) * np.pi, 0.0)

    def circuit (self, X:np.ndarray):
         qml.templates.AmplitudeEmbedding(X, wires = self.qubits[self.input_qubit_register], pad_with = 1.0, normalize = True)
         for i in range(self.n_layers):
             self.layer(layer_index = i)
         return [qml.expval(qml.PauliZ(qubit)) for qubit in self.qubits[self.ancilla_qubit_register]]

    def cost_function (self, y:np.ndarray, yhat:np.ndarray):
        return np.sum(np.subtract(y, yhat) ** 2) / y.shape[0]

    def train(self, X:np.ndarray, y:np.ndarray, batch_size:int=10, epochs:int=1):
        self.qnode = qml.QNode(self.circuit, device = self.device, diff_method = self.differentiation_method)
        for epoch in range(epochs):
            for i in range((X.shape[0] // batch_size) + 1):
                batch_index = np.random.randint(0, X.shape[0], (batch_size,))
                X_batch, y_batch = X[batch_index], y[batch_index]
                yhat_batch = np.array([self.qnode(X_batch[j,:], shots = self.shots) for j in range(batch_size)], dtype = "float32")
                self.theta = self.optimizer.step(lambda theta: self.cost_function(y_batch, yhat_batch), self.theta)

A Bit More Context: This issue seems to have arisen from a separate issue with the qnode decorator. Specifically, if you apply this decorator to a class method and provide kwdargs from the class instance (e.g. @qml.qnode(self.device)), then the decorator doesn’t work (from my understanding, this is expected and can be viewed as a design constraint imposed by the decorator approach).

An obvious first attempt at a workaround is to initialize the qnode within a class method that receives self as the first parameter as I did above:

    def train(self, X:np.ndarray, y:np.ndarray, batch_size:int=10, epochs:int=1):
        self.qnode = qml.QNode(self.circuit, device = self.device, diff_method = self.diff_method)
        # ...

From what I’ve read in the Pennylane documentation, it seems this isn’t the preferred method of instantiating a QNode, but for my use case it appears necessary given the challenge with decorators mentioned above.

I know that decorators are common to other JIT and ML libraries that consider specific hardware devices (numba, jax, tensorflow, pytorch, etc.). However, it seems to me that if the end-user wants to automate, say, many experiments with differentiation method as a variable, this library architecture inhibits their ability to build extensible classes on top of the library.

Any advice/guidance on how to handle this (or a response showing me the error of my ways!) would be very much appreciated.

Thanks,
Ben

P.S. When I output the tensor instantiated in the __init__ method, it does appear to be able to attach to autograd:

(
    tensor(
        [[[ 0.12316359, -0.03912076, -0.06289622],
        [ 0.63608486, -0.26805952, -0.00448637],
        [ 0.06343109, -0.25281549, -0.45259441],
        [ 0.08460864,  0.56978016,  0.12710347],
        [ 0.02829714,  0.02599053, -0.08118433],
        [ 0.37180194, -0.15878937, -0.42624673]]], 
    requires_grad=True), 
    0.0
)

Hi @Ben, welcome to the forum!

You could try ensuring the parameters are a differentiable tensor by using the PennyLane native numpy array. Here’s an example of how to do this.

from pennylane import numpy as np 
np.array([0.1, 0.2])

Output: tensor([0.1, 0.2], requires_grad=True)

Let me know if this works. If not, it would be useful to see a traceback of the issue.

Hi @CatalinaAlbornoz,

Thanks for the quick response! Unfortunately that is the version of NumPy I have been using. I’m also not getting a traceback – this appears to be a logical error but I can’t quite see where. Here’s a longer version of the code:

import sys, itertools, re, csv
import pennylane as qml
from pennylane import numpy as np

class QuantumBinaryClassifier(object):

	def __init__(self, model_config: dict, debug: bool = False):

		# Hyperparameters
		self.gamma 						= model_config["gamma"]

		# Architecture
		self.n_layers                   = model_config["n_layers"]
		self.m_qubits                   = model_config["m_qubits"]
		self.qubits                     = list(range(self.m_qubits))
		self.ancilla_qubit_register 	= slice(*model_config["ancilla_qubit_register"])
		self.input_qubit_register		= slice(*model_config["input_qubit_register"])
		self.theta_shape 				= (self.n_layers, self.m_qubits, 3)
		self.theta                      = (np.random.randn(*self.theta_shape) * np.pi * self.gamma, 0.0)

		# Hardware
		self.qnode 						= None
		self.device                     = qml.device(model_config["device"], wires = self.qubits)
		self.shots                      = model_config["shots"]

		# Data
		self.X                          = np.array([[]], dtype = "float32")
		self.y                          = np.array([], dtype = "float32")
		self.yhat                       = []

		# Optimization
		optimizer, optimizer_kwdargs    = model_config["optimizer"]
		self.optimizer                  = getattr(qml.optimize, optimizer)(**optimizer_kwdargs)
		self.differentiation_method 	= model_config["differentiation_method"]
		self.interface 					= model_config["interface"]

		# Developer
		self.debug                      = debug

		return None=

	def init_qnode (self):
		self.qnode = qml.QNode(self.circuit, device = self.device, interface = self.interface, diff_method = self.differentiation_method)
		return self

	# Example PQC
	def circuit (self, X):
		qml.templates.AmplitudeEmbedding(X, wires = self.qubits[self.input_qubit_register], pad_with = 1.0, normalize = True)
		params, _ = self.theta
		for layer_index in range(self.n_layers):
			for qubit in self.qubits[self.input_qubit_register]:
				qml.Hadamard(qubit)
			for qubit in self.qubits[self.input_qubit_register]:
				qml.Rot(self.theta[0][layer_index, qubit, 0], params[layer_index, qubit, 1], params[layer_index, qubit, 2], wires = qubit)
			for qubit in self.qubits[self.input_qubit_register]:
				qml.CNOT(wires = [qubit, (qubit + 1) % self.m_qubits])
		return [qml.expval(qml.PauliZ(qubit)) for qubit in self.qubits[self.ancilla_qubit_register]]

	def train (self, X, y, batch_size = 32, epochs = 10, draw_circuit = False, prediction_lambda = lambda x: np.where(x > 0.5, 1.0, 0.0)):
		self.X, self.y = X, y

		# Init QNode
		self.init_qnode()

		# Circuit drawing
		if draw_circuit:
			print(qml.draw(self.qnode, show_all_wires = True)(X[0,:]))

		# Training
		for epoch in range(epochs):

			# Get minibatch
			for i in range((X.shape[0] // batch_size) + 1):
				idx = np.random.randint(0, X.shape[0], (batch_size,))
				X_batch, y_batch = X[idx], y[idx]

				# Classify examples
				yhat_batch = np.array([self.qnode(X_batch[j,:], shots = self.shots) for j in range(batch_size)], dtype = "float32")

				# Compute cost, parameter updates, and predictions
				self.theta = self.optimizer.step(lambda theta: self.cost_function(theta, y_batch, yhat_batch), self.theta)
				print(self.theta[0][0,0,0])
				self.loss = self.cost_function(self.theta, y_batch, yhat_batch)
				batch_predictions = self.prediction_function(yhat_batch)

			self.yhat = np.array([self.qnode(X[i,:], shots = self.shots) for i in range(X.shape[0])], dtype = "float32").flatten()
			self.predictions = self.prediction_function(self.yhat)
			print(f"Epoch: {epoch + 1} / Training Accuracy: {self.accuracy(y, self.predictions):0.2f}% / n = {X.shape[0]}")

		return self

	def cost_function (self, theta, y, yhat):
		self.theta = theta
		return np.sum(np.subtract(y, yhat) ** 2) / y.shape[0]

	def prediction_function (self, yhat, threshold = 0.5):
		return np.where(yhat > threshold, 1.0, 0.0)

	def accuracy (self, targets, predictions):
		return 100 * np.sum(np.where(np.add(targets, predictions) % 2 == 0, 1.0, 0.0)) / targets.shape[0]

if __name__ == "__main__":
# Configure QML Model
	model_config = {
		"n_layers": 1,
		"m_qubits": 6,
		"device": "default.qubit",
		"input_qubit_register": (0,5),
		"ancilla_qubit_register": (5,6),
		"shots": 1024,
		"optimizer": ("NesterovMomentumOptimizer", {"stepsize" : 0.2, "momentum": 0.9}),
		"interface": "autograd",
		"differentiation_method": "parameter-shift",
		"gamma": 0.1
	}

	vqc = QuantumClassifier(
		model_config,
		debug = False
	).train(
		X_train,
		y_train,
		epochs = 2,
		draw_circuit = True,
	)

else:
    pass

Hi @Ben, I’m not being able to run your code. Could you please share the working Jupyter notebook? You can upload it here on a comment.

Thanks!

Hi @CatalinaAlbornoz, Apologies on not sharing a runnable example – I’ll make sure to do that next time. Turns out my references to the QNode class, instantiation, and circuit function were a bit off which caused issues when I passed the cost function in to optimizer.step(). I managed to get it sorted and its work now : )

If it helps, I can edit this (and the other posts) to show a working example of the error and how I resolved it. Let me know.

Thanks!

Hi @Ben, I’m glad you could solve it!

Yes, I think it would be great if you could add a comment on this thread showing your solution. It will probably help others in the future!