{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Differentiating CVQNN layers with Piquasso and Tensorflow"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Photonic variational quantum circuits (VQCs) are the analogs of traditional qubit-based VQCs adapted to photonic setting. A photonic VQC consists of the following:\n",
"\n",
"1. An initial state $\\ket{\\psi(\\mathbf{x})}$ (optionally) depending on some data $\\mathbf{x}$.\n",
"2. A photonic quantum circuit $\\hat{U}(\\mathbf{\\theta})$, which is parametrized by a collection of free parameters $\\mathbf{\\theta}$.\n",
"3. A set of observables $\\{ \\hat{O}_j \\}_{j=1}^N$ to be measured.\n",
"\n",
"In this setup, the expectation value of an observable $\\hat{O}_j$ is a function of the circuit parameters $\\mathbf{\\theta}$, i.e.,\n",
"$$\n",
" f_j(\\mathbf{\\theta}) := \\bra{\\psi_0} \\hat{U}^\\dagger(\\mathbf{\\theta}) \\hat{O}_j \\hat{U}(\\mathbf{\\theta}) \\ket{\\psi_0}.\n",
"$$\n",
"Typically a loss function is a function of the expectation values of these observables (or estimates thereof), and the training of a VQC amounts to tuning the circuit parameters $\\mathbf{x}$ to minimize the loss function.\n",
"\n",
"The increased success of solving difficult problems by deep learning in classical computing led to similar developments in variational quantum computing . In particular, a family of photonic VQCs called *continuous-variable quantum neural networks* (CVQNNs) can be considered as analogs of deep neural networks .\n",
"\n",
"The construction of CVQNNs is based on their classical counterpart. In classical deep learning, a fully connected neural network layer is often formulated as\n",
"$$\n",
" \\mathcal{L}(\\mathbf{x}) = \\varphi(W\\mathbf{x} + \\mathbf{b}),\n",
"$$\n",
"where $\\mathbf{x} \\in \\mathbb{R}^n,\\,\\mathbf{b}\\ \\in \\mathbb{R}^m$ are the input and the bias vectors, respectively, $W \\in \\mathbb{R}^{m \\times n}$ is the matrix containing the weights of the layer, and $\\varphi$ is a non-linear activation function, usually a sigmoid, ReLU or hyperbolic tangent function.\n",
"\n",
"Building upon this, the CVQNN construction admits layers of linear optical gates, which can implement any symplectic affine transformations on the phase space . Subsequently, a layer of nonlinear optical gates corresponds to an activation function.\n",
"Hence, a single CVQNN layer depending on weights $\\mathbf{x}$ as gate parameters can be written as\n",
"$$\n",
" \\mathcal{L}(\\mathbf{x}) :=\n",
" \\Phi(\\mathbf{x}_{\\Phi})\\, D(\\mathbf{x}_{D})\\, I_2(\\mathbf{x}_{I_2}) \\, S(\\mathbf{x}_{S}) \\, I_1(\\mathbf{x}_{I_1}),\n",
"$$\n",
"where $\\Phi$ represents a column of non-linear optical gates, $I_i$ represent [interferometers](../instructions/gates.rst#piquasso.instructions.gates.Interferometer), and\n",
"$$\n",
" D(\\mathbf{\\alpha}) = \\bigotimes_{i=1}^d D_i(\\alpha_i) \\qquad (\\mathbf{\\alpha} \\in \\mathbb{C}^d), \\\\\n",
" S(\\mathbf{z}) = \\bigotimes_{i=1}^d S_i(z_i) \\qquad (\\mathbf{z} \\in \\mathbb{C}^d)\n",
"$$\n",
"are columns of [Displacement](../instructions/gates.rst#piquasso.instructions.gates.Displacement) and [Squeezing](../instructions/gates.rst#piquasso.instructions.gates.Displacement) gates, respectively.\n",
"Usually, the $\\Phi$ is chosen to be a column of [Kerr](../instructions/gates.rst#piquasso.instructions.gates.Kerr) gates, i.e.,\n",
"$$\n",
" \\Phi(\\mathbf{x}_\\Phi) := K(\\mathbf{\\kappa}) = \\bigotimes_{i=1}^d K_i(\\kappa_i) \\qquad (\\kappa_i \\in \\mathbb{R}).\n",
"$$\n",
"The parameters $\\mathbf{x}_{\\Phi}, \\mathbf{x}_{I_2}, \\mathbf{x}_{I_1}$ are real-valued, whereas $\\mathbf{x}_{D}$ and $\\mathbf{x}_{S}$ are complex-valued.\n",
"Interferometers can be further decomposed into phaseshifter and beamsplitter gates using the Clements decomposition which yields a parametrization in terms of phaseshifter and beamsplitter angles.\n",
"\n",
"In parallel with the expression for a fully connected layer in the classical setting, one can also get a similar expression for a CVQNN layer as follows. Considering the effect of linear optical gates on the ladder operators $\\mathbf{\\xi}$, the layer can be written as\n",
"$$\n",
" \\mathcal{L}^\\dagger(\\mathbf{x}) \\mathbf{\\xi} \\mathcal{L}(\\mathbf{x}) = K^\\dagger( \\mathbf{x}_{\\text{non-linear}}) (S_{(c)} \\mathbf{\\xi} + \\mathbf{\\alpha} I) K( \\mathbf{x}_{\\text{non-linear}}),\n",
"$$\n",
"where $\\mathbf{x} = (\\mathbf{x}_{\\text{non-linear}}, \\mathbf{x}_{\\text{linear}}) = (\\mathbf{x}_{\\text{non-linear}}, \\mathbf{x}_{D}, \\mathbf{x}_{I_2}, \\mathbf{x}_{S}, \\mathbf{x}_{I_1})$\n",
"and $S_{(c)}, \\mathbf{\\alpha}$ depend on $\\mathbf{x}_{\\text{linear}}$. Here, the term $S_{(c)} \\mathbf{\\xi} + \\mathbf{\\alpha} I$ represents a symplectic affine transformation on the ladder operator $\\mathbf{\\xi}$, and conjugation by the Kerr gates represent a non-linear transformation."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In Piquasso, one can easily create CVQNN circuits using the [Piquasso CVQNN](../advanced/cvqnn.rst) module."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"import piquasso as pq\n",
"\n",
"d = 2 # Number of qumodes\n",
"\n",
"layer_count = 10 # Number of CVQNN layers\n",
"\n",
"# Generating random weights\n",
"weights = pq.cvqnn.generate_random_cvqnn_weights(layer_count=layer_count, d=d)\n",
"\n",
"# Creating CVQNN layers as a Piquasso subprogram\n",
"cvqnn_layers = pq.cvqnn.create_layers(weights)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Piquasso automatically sets up a subprogram containing the instructions of the desired CVQNN layer. Now we can embed this subprogram in any Piquasso program. Let's choose the input state as a pure displaced state."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"# The program definition\n",
"with pq.Program() as program:\n",
" pq.Q() | pq.Vacuum()\n",
"\n",
" for i in range(d):\n",
" pq.Q(i) | pq.Displacement(r=0.1)\n",
"\n",
" pq.Q() | cvqnn_layers"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we need to choose the simulator which executes the instructions. Since a CVQNN layer includes non-linear terms, we definitely need to perform the simulation in Fock space. Since our initial state is pure, we can use [PureFockSimulator](../simulators/fock.rst#piquasso._simulators.fock.pure.simulator.PureFockSimulator)."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"cutoff = 5\n",
"\n",
"simulator = pq.PureFockSimulator(d, pq.Config(cutoff=cutoff))\n",
"\n",
"final_state = simulator.execute(program).state"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After obtaining the state, we can calculate several things, e.g. the expectation value of the position operator on mode 0."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Mean position on mode 0:\n"
]
},
{
"data": {
"text/plain": [
"0.08810452589019563"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"print(\"Mean position on mode 0:\")\n",
"final_state.mean_position(mode=0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In order to differentiate quantitities, we need to modify the simulation. In itself, `PureFockSimulator` is unable to perform automatic differentiation. In order to do that, we can use `TensorflowConnector`, which replaces NumPy to TensorFlow under the hood. For a concrete example, let the loss function be\n",
"$$\n",
"J(w) = \\| \\ket{\\psi(w)} - \\ket{\\psi_*} \\|_2,\n",
"$$\n",
"where $\\ket{\\psi(w)}$ is the final state of the circuit, and $\\ket{\\psi_*}$ is some random final state."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"from scipy.special import comb\n",
"\n",
"state_vector_size = comb(d + cutoff - 1, cutoff - 1, exact=True)\n",
"\n",
"psi_star = np.random.rand(state_vector_size) + 1j * np.random.rand(state_vector_size)\n",
"\n",
"psi_star /= np.sum(np.abs(psi_star))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then, by using [tf.GradientTape](https://www.tensorflow.org/api_docs/python/tf/GradientTape), we can differentiate this loss function."
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loss: 1.9604476480647497\n",
"Loss gradient: [[ 2.48370504e-02 6.91253512e-03 -5.17538228e-02 5.91197083e-01\n",
" 8.35270137e-01 4.33414809e-02 -1.44455276e-03 -6.10485008e-02\n",
" 6.30518617e-01 7.87589392e-01 4.15688561e-04 6.45192857e-03\n",
" -7.05893753e-02 -1.81552492e-02]\n",
" [ 5.51070724e-02 1.14482055e-02 -7.20810177e-02 8.70276885e-01\n",
" 6.50958369e-01 4.79129674e-02 -1.05355280e-02 -6.39677085e-02\n",
" 7.03435046e-01 7.72832475e-01 6.62571242e-03 1.03948801e-03\n",
" -7.14190194e-02 -1.32908739e-02]\n",
" [ 5.91775056e-02 4.75369308e-03 -6.20956892e-02 8.06572060e-01\n",
" 7.06752388e-01 4.33605125e-02 -4.52911124e-04 -6.26476669e-02\n",
" 7.98818963e-01 7.21751009e-01 1.01077259e-02 -2.26709094e-03\n",
" -6.75334937e-02 -1.38186903e-02]\n",
" [ 4.01203598e-02 1.99768691e-03 -5.45376280e-02 8.81982242e-01\n",
" 6.53387621e-01 5.93868412e-03 -6.30429008e-04 -5.38798294e-02\n",
" 9.05422584e-01 7.15949040e-01 -3.64702223e-04 -5.77410545e-05\n",
" -6.82197114e-02 1.74116637e-03]\n",
" [ 1.07959822e-03 1.73058480e-04 -5.44175901e-02 8.66129564e-01\n",
" 6.70398558e-01 1.66385446e-02 4.30945595e-04 -5.41922390e-02\n",
" 9.00744709e-01 7.08804040e-01 4.50634059e-03 5.91755415e-04\n",
" -6.23273595e-02 5.69621412e-03]\n",
" [ 4.21102694e-03 3.93742722e-03 -5.36233257e-02 1.03238139e+00\n",
" 5.26661099e-01 -1.52360597e-03 -3.54072884e-03 -4.91789820e-02\n",
" 9.16184120e-01 6.90672795e-01 1.16414634e-03 -2.31269771e-03\n",
" -5.94908198e-02 9.16505767e-03]\n",
" [-2.04315766e-02 5.76912708e-04 -4.85917483e-02 8.56742917e-01\n",
" 6.53319067e-01 -1.99088351e-02 8.87190361e-04 -4.30689726e-02\n",
" 9.27679869e-01 6.49967198e-01 1.48523330e-04 -1.12318622e-03\n",
" -4.57634914e-02 1.02198966e-02]\n",
" [-2.21043277e-02 1.45639573e-03 -4.43768450e-02 9.80622059e-01\n",
" 5.40051984e-01 -4.77245102e-02 2.81496583e-03 -6.23762968e-02\n",
" 1.05210311e+00 4.76482110e-01 -8.71547209e-04 6.93330613e-05\n",
" -8.64155789e-02 8.37255149e-03]\n",
" [-4.67122038e-02 -3.90292564e-05 -6.32088147e-02 1.12961323e+00\n",
" 4.26270074e-01 -2.25618686e-02 -5.57393577e-03 -5.11763569e-02\n",
" 9.44769422e-01 6.50075251e-01 -1.20220424e-03 -2.19911083e-03\n",
" -6.92235563e-02 6.31057580e-03]\n",
" [-3.85113231e-02 -1.01100459e-03 -5.13675566e-02 7.89344621e-01\n",
" 6.40349720e-01 -3.04770202e-02 -4.07158955e-03 -4.26573559e-02\n",
" 8.30388031e-01 7.82317691e-01 2.13340682e-03 1.06722282e-03\n",
" -4.83288207e-02 5.57853420e-03]]\n"
]
}
],
"source": [
"import tensorflow as tf\n",
"\n",
"tf.get_logger().setLevel(\"ERROR\") # Turns off complex->float casting warnings\n",
"\n",
"simulator = pq.PureFockSimulator(\n",
" d, pq.Config(cutoff=cutoff), connector=pq.TensorflowConnector()\n",
")\n",
"\n",
"w = tf.Variable(weights)\n",
"psi_star = tf.Variable(psi_star)\n",
"\n",
"with tf.GradientTape() as tape:\n",
" cvqnn_layers = pq.cvqnn.create_layers(w)\n",
"\n",
" with pq.Program() as program:\n",
" pq.Q() | pq.Vacuum()\n",
"\n",
" for i in range(d):\n",
" pq.Q(i) | pq.Displacement(r=0.1)\n",
"\n",
" pq.Q() | cvqnn_layers\n",
"\n",
" simulator.execute(program)\n",
"\n",
" final_state = simulator.execute(program).state\n",
"\n",
" psi = final_state.state_vector\n",
"\n",
" loss = tf.math.reduce_sum(tf.math.abs(psi - psi_star))\n",
"\n",
"loss_grad = tape.gradient(loss, w)\n",
"\n",
"print(f\"Loss: {loss.numpy()}\")\n",
"print(f\"Loss gradient: {loss_grad.numpy()}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"However, Piquasso is written in a way that it supports `tf.function` (see [Better performance with tf.function](https://www.tensorflow.org/guide/function)) one can also use `tf.function` for this task. Refactoring everything into a function, we can use the `tf.function` decorator. Note, that we have to avoid side effects in any function decorated with `tf.function`, because side effects are only executed at the tracing step. Therefore, instantiation of `pq.Program` should happen by providing the instructions as constructor arguments, instead of using the `with` keyword."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loss: 2.018569979174099\n",
"Loss gradient: [[ 0.00000000e+00 0.00000000e+00 0.00000000e+00 5.75160857e-01\n",
" 6.83322458e-01 1.28450675e-02 -4.00008041e-05 -6.30548348e-03\n",
" -7.49845629e-01 -8.00313533e-01 7.09929231e-05 1.03634563e-02\n",
" -1.18960333e-02 1.04187002e-02]\n",
" [-1.35800236e-03 -5.49585214e-04 -5.68490534e-03 7.68436902e-01\n",
" 5.05490527e-01 6.69074565e-03 6.61089567e-04 -5.98396524e-03\n",
" -7.70615878e-01 -7.65752853e-01 2.44147270e-03 2.03468039e-03\n",
" -4.38080826e-03 1.26999161e-02]\n",
" [ 1.33582634e-02 -1.68655402e-03 -1.85593853e-03 7.07517913e-01\n",
" 5.44470265e-01 2.49539746e-02 6.42390173e-05 -1.95025626e-03\n",
" -8.25024061e-01 -6.44185279e-01 6.53653384e-03 -7.65527162e-03\n",
" 9.10541653e-03 3.40278572e-03]\n",
" [ 4.26398634e-02 -1.80787941e-03 6.39415698e-03 7.10539771e-01\n",
" 5.16317806e-01 4.19027819e-02 3.11151157e-04 6.34437434e-03\n",
" -7.11899830e-01 -5.86043987e-01 -1.78246076e-03 -3.32723452e-04\n",
" 1.04650909e-02 4.07298677e-03]\n",
" [ 4.29666003e-02 -5.13803858e-05 4.61329396e-03 6.98424383e-01\n",
" 5.25555203e-01 3.78096485e-02 -3.26929949e-04 6.85206072e-03\n",
" -7.86593321e-01 -5.59964717e-01 8.99000290e-03 4.50767474e-03\n",
" 2.76895247e-02 7.51784650e-03]\n",
" [ 4.84763174e-02 -5.06656057e-03 2.09086242e-02 7.78957857e-01\n",
" 4.11735915e-01 4.47041965e-02 5.57595820e-03 1.66618932e-02\n",
" -7.45897045e-01 -6.70781724e-01 4.18901129e-03 -5.51458453e-03\n",
" 3.53793166e-02 2.84834471e-03]\n",
" [ 6.30724855e-02 -6.41540211e-04 2.14924447e-02 6.11241592e-01\n",
" 5.05390366e-01 6.89479439e-02 -1.41398304e-03 2.94560107e-02\n",
" -8.62929995e-01 -6.09878053e-01 1.88758840e-04 -3.11469367e-03\n",
" 5.47604296e-02 -3.52522121e-03]\n",
" [ 6.68569098e-02 -2.65481869e-03 3.22995882e-02 7.03713041e-01\n",
" 4.32114569e-01 4.74536273e-02 -3.67396656e-03 2.02959096e-02\n",
" -8.42787956e-01 -4.25121040e-01 -3.99655314e-03 2.98835285e-04\n",
" 1.63926716e-02 -1.21762767e-02]\n",
" [ 4.79882382e-02 6.12290682e-05 1.62381274e-02 8.37636878e-01\n",
" 3.50979911e-01 5.30524902e-02 7.52035617e-03 1.63628512e-02\n",
" -8.23920872e-01 -6.25385991e-01 -2.62846126e-03 -5.16820365e-03\n",
" 2.70939337e-02 -7.64525657e-03]\n",
" [ 6.78217682e-02 5.96498824e-04 1.31378911e-02 5.59849745e-01\n",
" 5.20217392e-01 7.60063612e-02 2.26111302e-03 1.50965198e-02\n",
" -4.33486232e-01 -6.76081337e-01 -5.44203449e-03 4.00677625e-03\n",
" 2.97522697e-02 1.53461233e-03]]\n"
]
}
],
"source": [
"def calculate_loss(w, psi_star, cutoff):\n",
" d = pq.cvqnn.get_number_of_modes(w.shape[1])\n",
"\n",
" simulator = pq.PureFockSimulator(\n",
" d,\n",
" pq.Config(cutoff=cutoff),\n",
" connector=pq.TensorflowConnector(decorate_with=tf.function),\n",
" )\n",
"\n",
" with tf.GradientTape() as tape:\n",
" cvqnn_layers = pq.cvqnn.create_layers(w)\n",
"\n",
" final_state = simulator.execute_instructions(\n",
" instructions=[pq.Vacuum()] + cvqnn_layers.instructions\n",
" ).state\n",
"\n",
" psi = final_state.state_vector\n",
"\n",
" loss = tf.math.reduce_sum(tf.math.abs(psi - psi_star))\n",
"\n",
" return loss, tape.gradient(loss, w)\n",
"\n",
"\n",
"improved_calculate_loss = tf.function(calculate_loss)\n",
"\n",
"loss, loss_grad = improved_calculate_loss(w, psi_star, cutoff)\n",
"\n",
"print(f\"Loss: {loss.numpy()}\")\n",
"print(f\"Loss gradient: {loss_grad.numpy()}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The first run is called the tracing step, and it takes some time, because Tensorflow captures a [tf.Graph](https://www.tensorflow.org/api_docs/python/tf/Graph) here. The size of the graph can be decreased by passing the `decorate_with=tf.function` argument to `pq.TensorflowConnector`, which also decreases the execution time of the tracing step. After the first run, a significant speedup is observed. We can also compare the runtimes of the compiled and non-compiled function."
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Regular: 2.4323306560516356 s (+/- 0.25203824522512575 s).\n",
"Improved: 0.011875104904174805 s (+/- 0.0002783859661023313 s).\n"
]
}
],
"source": [
"import time\n",
"import numpy as np\n",
"\n",
"regular_runtimes = []\n",
"improved_runtimes = []\n",
"\n",
"for i in range(10):\n",
" w = tf.Variable(pq.cvqnn.generate_random_cvqnn_weights(layer_count, d))\n",
"\n",
" start_time = time.time()\n",
" calculate_loss(w, psi_star, cutoff)\n",
" end_time = time.time()\n",
"\n",
" regular_runtimes.append(end_time - start_time)\n",
"\n",
" start_time = time.time()\n",
" improved_calculate_loss(w, psi_star, cutoff)\n",
" end_time = time.time()\n",
"\n",
" improved_runtimes.append(end_time - start_time)\n",
"\n",
"print(f\"Regular: {np.mean(regular_runtimes)} s (+/- {np.std(regular_runtimes)} s).\")\n",
"print(f\"Improved: {np.mean(improved_runtimes)} s (+/- {np.std(improved_runtimes)} s).\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"But that is not everything yet! One can also create a similar function with the `jit_compile=True` flag, since every operation in Piquasso can be JIT-compiled using XLA through `tf.function`."
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2024-03-02 09:40:10.970536: I external/local_xla/xla/service/service.cc:168] XLA service 0x23ff2f70 initialized for platform Host (this does not guarantee that XLA will be used). Devices:\n",
"2024-03-02 09:40:10.970561: I external/local_xla/xla/service/service.cc:176] StreamExecutor device (0): Host, Default Version\n",
"2024-03-02 09:40:14.366437: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.\n",
"2024-03-02 09:42:47.120088: E external/local_xla/xla/service/slow_operation_alarm.cc:65] \n",
"********************************\n",
"[Compiling module a_inference_calculate_loss_184206__XlaMustCompile_true_config_proto_3175580994766145631_executor_type_11160318154034397263_.64111] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.\n",
"********************************\n",
"2024-03-02 09:45:40.915775: E external/local_xla/xla/service/slow_operation_alarm.cc:133] The operation took 4m53.795791141s\n",
"\n",
"********************************\n",
"[Compiling module a_inference_calculate_loss_184206__XlaMustCompile_true_config_proto_3175580994766145631_executor_type_11160318154034397263_.64111] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.\n",
"********************************\n",
"WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
"I0000 00:00:1709394340.917125 34412 device_compiler.h:186] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process.\n",
"2024-03-02 09:45:41.016959: E external/local_xla/xla/stream_executor/stream_executor_internal.h:177] SetPriority unimplemented for this stream.\n"
]
},
{
"data": {
"text/plain": [
"(,\n",
" )"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"jit_compiled_calculate_loss = tf.function(jit_compile=True)(calculate_loss)\n",
"\n",
"jit_compiled_calculate_loss(w, psi_star, cutoff)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Compiling the same function takes significantly more time, but the compilation step results in an extra order of magnitude runtime improvement:"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Regular:\t2.4323306560516356 s (+/- 0.25203824522512575 s).\n",
"Improved:\t0.011875104904174805 s (+/- 0.0002783859661023313 s).\n",
"JIT compiled:\t0.0013802051544189453 s (+/- 0.00014840260989162237 s).\n"
]
}
],
"source": [
"jit_compiled_runtimes = []\n",
"\n",
"for i in range(10):\n",
" w = tf.Variable(pq.cvqnn.generate_random_cvqnn_weights(layer_count, d))\n",
"\n",
" start_time = time.time()\n",
" jit_compiled_calculate_loss(w, psi_star, cutoff)\n",
" end_time = time.time()\n",
"\n",
" jit_compiled_runtimes.append(end_time - start_time)\n",
"\n",
"print(f\"Regular:\\t{np.mean(regular_runtimes)} s (+/- {np.std(regular_runtimes)} s).\")\n",
"print(f\"Improved:\\t{np.mean(improved_runtimes)} s (+/- {np.std(improved_runtimes)} s).\")\n",
"print(\n",
" f\"JIT compiled:\\t{np.mean(jit_compiled_runtimes)} s \"\n",
" f\"(+/- {np.std(jit_compiled_runtimes)} s).\"\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We are able to use this function for. e.g., quantum state learning. Consider the state\n",
"$$\n",
"\\ket{\\psi_*} = \\frac{1}{\\sqrt{2}} \\left ( \\ket{03} + \\ket{30} \\right ),\n",
"$$\n",
"which is an example of a so-called NOON state. We can produce this using Piquasso:"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [],
"source": [
"with pq.Program() as target_state_preparation:\n",
" pq.Q(all) | pq.StateVector([0, 3]) / np.sqrt(2)\n",
" pq.Q(all) | pq.StateVector([3, 0]) / np.sqrt(2)\n",
"\n",
"\n",
"target_state = simulator.execute(target_state_preparation).state\n",
"\n",
"target_state_vector = target_state.state_vector\n",
"\n",
"psi_star = tf.Variable(target_state_vector)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we can demonstrate the speed of the JIT-compiled calculation by creating a simple optimization algorithm as follows:"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"500:\t\t 0.8607086996549679\n",
"1000:\t\t 0.8473090910243798\n",
"1500:\t\t 0.5849377384183669\n",
"2000:\t\t 0.40879991260231896\n",
"2500:\t\t 0.29813437623431027\n",
"3000:\t\t 0.2761913995731654\n",
"3500:\t\t 0.26621822690380204\n",
"4000:\t\t 0.23435492471006458\n",
"4500:\t\t 0.23874428115482763\n",
"5000:\t\t 0.2205772331178741\n",
"5500:\t\t 0.2416911247101965\n",
"6000:\t\t 0.2146732713315681\n",
"6500:\t\t 0.1993502777279203\n",
"7000:\t\t 0.22721179422812682\n",
"7500:\t\t 0.21320290731217184\n",
"8000:\t\t 0.20555180951647123\n",
"8500:\t\t 0.21018212032348688\n",
"9000:\t\t 0.20439647783148554\n",
"9500:\t\t 0.19759860554592923\n",
"10000:\t\t 0.18922056395791387\n",
"Final loss:\t 0.18922056395791387\n"
]
}
],
"source": [
"learning_rate = 0.01\n",
"iterations = 10000\n",
"\n",
"optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
"\n",
"w = tf.Variable(pq.cvqnn.generate_random_cvqnn_weights(layer_count, d))\n",
"\n",
"\n",
"for i in range(iterations):\n",
" loss, loss_grad = jit_compiled_calculate_loss(w, psi_star, cutoff)\n",
"\n",
" optimizer.apply_gradients(zip([loss_grad], [w]))\n",
"\n",
" if (i + 1) % (iterations // 20) == 0:\n",
" print(f\"{i + 1}:\\t\\t\", loss.numpy())\n",
"\n",
"print(\"Final loss:\\t\", loss.numpy())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can use the final weights to calculate the final state, and calculate its fidelity with the target state."
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Final state fidelity: 0.9994946525680765\n"
]
}
],
"source": [
"program = pq.cvqnn.create_program(w)\n",
"\n",
"final_state = simulator.execute(program).state\n",
"\n",
"print(\"Final state fidelity: \", final_state.fidelity(target_state))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.16"
}
},
"nbformat": 4,
"nbformat_minor": 2
}