Not numpy from pennylane getting started (#932)

In this PR I will replace "numpy from pennylane" with other
alternatives.
(I had a meeting with Tom to decide what to use in each case)

---------

Co-authored-by: Tom Bromley <[email protected]>

demonstrations/adjoint_diff_benchmarking.py

@@ -22,19 +22,22 @@
 import timeit
 import matplotlib.pyplot as plt
 import pennylane as qml
-from pennylane import numpy as np
+import jax
+
+jax.config.update("jax_platform_name", "cpu")
 
 plt.style.use("bmh")
 
 n_samples = 5
 
+
 def get_time(qnode, params):
-    globals_dict = {'grad': qml.grad, 'circuit': qnode, 'params': params}
+    globals_dict = {'grad': jax.grad, 'circuit': qnode, 'params': params}
     return timeit.timeit("grad(circuit)(params)", globals=globals_dict, number=n_samples)
 
-def wires_scaling(n_wires, n_layers):
 
-    rng = np.random.default_rng(seed=42)
+def wires_scaling(n_wires, n_layers):
+    key = jax.random.PRNGKey(42)
 
     t_adjoint = []
     t_ps = []
@@ -45,17 +48,16 @@ def circuit(params, wires):
         return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1) @ qml.PauliZ(2))
 
     for i_wires in n_wires:
-
         dev = qml.device("lightning.qubit", wires=i_wires)
         dev_python = qml.device("default.qubit", wires=i_wires)
-        
+
         circuit_adjoint = qml.QNode(lambda x: circuit(x, wires=i_wires), dev, diff_method="adjoint")
         circuit_ps = qml.QNode(lambda x: circuit(x, wires=i_wires), dev, diff_method="parameter-shift")
         circuit_backprop = qml.QNode(lambda x: circuit(x, wires=i_wires), dev_python, diff_method="backprop")
 
         # set up the parameters
         param_shape = qml.StronglyEntanglingLayers.shape(n_wires=i_wires, n_layers=n_layers)
-        params = rng.standard_normal(param_shape, requires_grad=True)
+        params = jax.random.normal(key, param_shape)
 
         t_adjoint.append(get_time(circuit_adjoint, params))
         t_backprop.append(get_time(circuit_backprop, params))
@@ -65,8 +67,7 @@ def circuit(params, wires):
 
 
 def layers_scaling(n_wires, n_layers):
-
-    rng = np.random.default_rng(seed=42)
+    key = jax.random.PRNGKey(42)
 
     dev = qml.device("lightning.qubit", wires=n_wires)
     dev_python = qml.device('default.qubit', wires=n_wires)
@@ -84,24 +85,22 @@ def circuit(params):
     circuit_backprop = qml.QNode(circuit, dev_python, diff_method="backprop")
 
     for i_layers in n_layers:
-
         # set up the parameters
         param_shape = qml.StronglyEntanglingLayers.shape(n_wires=n_wires, n_layers=i_layers)
-        params = rng.standard_normal(param_shape, requires_grad=True)
-        
+        params = jax.random.normal(key, param_shape)
+
         t_adjoint.append(get_time(circuit_adjoint, params))
         t_backprop.append(get_time(circuit_backprop, params))
         t_ps.append(get_time(circuit_ps, params))
 
     return t_adjoint, t_backprop, t_ps
 
-if __name__ == "__main__":
 
-    layers_list= [1, 3, 5, 7, 9]
+if __name__ == "__main__":
+    layers_list = [1, 3, 5, 7, 9]
     n_wires = 5
     adjoint_layers, backprop_layers, ps_layers = layers_scaling(n_wires, layers_list)
 
-
     wires_list = [3, 6, 9, 12, 15]
     n_layers = 3
     adjoint_wires, backprop_wires, ps_wires = wires_scaling(wires_list, n_layers)

demonstrations/tutorial_adjoint_diff.py

@@ -68,10 +68,14 @@
 # So how does it work? Instead of jumping straight to the algorithm, let's explore the above equations
 # and their implementation in a bit more detail.
 #
-# To start, we import PennyLane and PennyLane's numpy:
+# To start, we import PennyLane and Jax's numpy:
 
 import pennylane as qml
-from pennylane import numpy as np
+import jax
+from jax import numpy as np
+
+jax.config.update("jax_platform_name", "cpu")
+
 
 ##############################################################################
 # We also need a circuit to simulate:
@@ -376,9 +380,9 @@ def circuit(a):
 # since we calculated them in reverse
 grads = grads[::-1]
 
-print("our calculation: ", grads)
+print("our calculation: ", [float(grad) for grad in grads])
 
-grad_compare = qml.grad(circuit)(x)
+grad_compare = jax.grad(circuit)(x)
 print("comparison: ", grad_compare)
 
 ##############################################################################
@@ -399,7 +403,7 @@ def circuit_adjoint(a):
     qml.RZ(a[2], wires=1)
     return qml.expval(M)
 
-print(qml.grad(circuit_adjoint)(x))
+print(jax.grad(circuit_adjoint)(x))
 
 ##############################################################################
 # Performance
@@ -411,8 +415,8 @@ def circuit_adjoint(a):
 # Backpropagation demonstrates decent time scaling, but requires more and more
 # memory as the circuit gets larger.  Simulation of large circuits is already
 # RAM-limited, and backpropagation constrains the size of possible circuits even more.
-# PennyLane also achieves backpropagation derivatives from a Python simulator and
-# interface-specific functions. The ``"lightning.qubit"`` device does not support
+# PennyLane also achieves backpropagation derivatives from a Python simulator and interface-specific functions.
+# The ``"lightning.qubit"`` device does not support
 # backpropagation, so backpropagation derivatives lose the speedup from an optimized
 # simulator.
 #

demonstrations/tutorial_backprop.py

@@ -55,14 +55,20 @@
 Let's have a go implementing the parameter-shift rule manually in PennyLane.
 """
 import pennylane as qml
-from pennylane import numpy as np
+from jax import numpy as np
 from matplotlib import pyplot as plt
+import jax
+
+jax.config.update("jax_platform_name", "cpu")
+
+# set the random seed
+key = jax.random.PRNGKey(42)
 
 
 # create a device to execute the circuit on
 dev = qml.device("default.qubit", wires=3)
 
-@qml.qnode(dev, diff_method="parameter-shift", interface="autograd")
+@qml.qnode(dev, diff_method="parameter-shift")
 def circuit(params):
     qml.RX(params[0], wires=0)
     qml.RY(params[1], wires=1)
@@ -82,8 +88,8 @@ def circuit(params):
 # Let's test the variational circuit evaluation with some parameter input:
 
 # initial parameters
-rng = np.random.default_rng(1949320)
-params = rng.random([6], requires_grad=True)
+params = jax.random.normal(key, [6])
+
 
 print("Parameters:", params)
 print("Expectation value:", circuit(params))
@@ -102,10 +108,10 @@ def circuit(params):
 
 def parameter_shift_term(qnode, params, i):
     shifted = params.copy()
-    shifted[i] += np.pi/2
+    shifted = shifted.at[i].add(np.pi/2)
     forward = qnode(shifted)  # forward evaluation
 
-    shifted[i] -= np.pi
+    shifted = shifted.at[i].add(-np.pi/2)
     backward = qnode(shifted) # backward evaluation
 
     return 0.5 * (forward - backward)
@@ -121,20 +127,19 @@ def parameter_shift(qnode, params):
     gradients = np.zeros([len(params)])
 
     for i in range(len(params)):
-        gradients[i] = parameter_shift_term(qnode, params, i)
+        gradients = gradients.at[i].set(parameter_shift_term(qnode, params, i))
 
     return gradients
 
 print(parameter_shift(circuit, params))
 
 ##############################################################################
 # We can compare this to PennyLane's *built-in* quantum gradient support by using
-# the :func:`qml.grad <pennylane.grad>` function, which allows us to compute gradients
-# of hybrid quantum-classical cost functions. Remember, when we defined the
+# the ``jax.grad`` function. Remember, when we defined the
 # QNode, we specified that we wanted it to be differentiable using the parameter-shift
 # method (``diff_method="parameter-shift"``).
 
-grad_function = qml.grad(circuit)
+grad_function = jax.grad(circuit)
 print(grad_function(params)[0])
 
 ##############################################################################
@@ -143,7 +148,6 @@ def parameter_shift(qnode, params):
 
 print(qml.gradients.param_shift(circuit)(params))
 
-
 ##############################################################################
 # If you count the number of quantum evaluations, you will notice that we had to evaluate the circuit
 # ``2*len(params)`` number of times in order to compute the quantum gradient with respect to all
@@ -166,14 +170,15 @@ def parameter_shift(qnode, params):
 
 dev = qml.device("default.qubit", wires=4)
 
-@qml.qnode(dev, diff_method="parameter-shift", interface="autograd")
+@qml.qnode(dev, diff_method="parameter-shift")
 def circuit(params):
     qml.StronglyEntanglingLayers(params, wires=[0, 1, 2, 3])
     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliZ(3))
 
 # initialize circuit parameters
 param_shape = qml.StronglyEntanglingLayers.shape(n_wires=4, n_layers=15)
-params = rng.normal(scale=0.1, size=param_shape, requires_grad=True)
+params = jax.random.normal(key, param_shape) * 0.1
+
 print(params.size)
 print(circuit(params))
 
@@ -196,7 +201,7 @@ def circuit(params):
 # and see how this compares.
 
 # create the gradient function
-grad_fn = qml.grad(circuit)
+grad_fn = jax.grad(circuit)
 
 times = timeit.repeat("grad_fn(params)", globals=globals(), number=num, repeat=reps)
 backward_time = min(times) / num
@@ -253,7 +258,7 @@ def circuit(params):
 # One such device is :class:`default.qubit <pennylane.devices.DefaultQubit>`. It
 # has backends written using TensorFlow, JAX, and Autograd, so when used with the
 # TensorFlow, JAX, and Autograd interfaces respectively, supports backpropagation.
-# In this demo, we will use the default Autograd interface.
+# In this demo, we will use the JAX interface.
 
 dev = qml.device("default.qubit", wires=4)
 
@@ -263,14 +268,15 @@ def circuit(params):
 # mode* for the ``default.qubit`` device.
 
 
-@qml.qnode(dev, diff_method="backprop", interface="autograd")
+@qml.qnode(dev, diff_method="backprop")
 def circuit(params):
     qml.StronglyEntanglingLayers(params, wires=[0, 1, 2, 3])
     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliZ(3))
 
 # initialize circuit parameters
 param_shape = qml.StronglyEntanglingLayers.shape(n_wires=4, n_layers=15)
-params = rng.normal(scale=0.1, size=param_shape, requires_grad=True)
+params = jax.random.normal(key, param_shape) * 0.1
+
 print(circuit(params))
 
 ##############################################################################
@@ -290,7 +296,7 @@ def circuit(params):
 # overhead from using backpropagation. We can now estimate the time required to perform a
 # gradient computation via backpropagation:
 
-times = timeit.repeat("qml.grad(circuit)(params)", globals=globals(), number=num, repeat=reps)
+times = timeit.repeat("jax.grad(circuit)(params)", globals=globals(), number=num, repeat=reps)
 backward_time = min(times) / num
 print(f"Backward pass (best of {reps}): {backward_time} sec per loop")
 
@@ -327,16 +333,21 @@ def circuit(params):
 forward_backprop = []
 gradient_backprop = []
 
+qnode_shift = jax.jit(qml.QNode(circuit, dev, diff_method="parameter-shift"))
+qnode_backprop = jax.jit(qml.QNode(circuit, dev, diff_method="backprop"))
+
+grad_qnode_shift = jax.jit(jax.grad(qml.QNode(circuit, dev, diff_method="parameter-shift")))
+grad_qnode_backprop = jax.jit(jax.grad(qml.QNode(circuit, dev, diff_method="backprop")))
+
 for depth in range(0, 21):
     param_shape = qml.StronglyEntanglingLayers.shape(n_wires=4, n_layers=depth)
-    params = rng.normal(scale=0.1, size=param_shape, requires_grad=True)
+    params = jax.random.normal(key, param_shape) * 0.1
+
     num_params = params.size
 
     # forward pass timing
     # ===================
 
-    qnode_shift = qml.QNode(circuit, dev, diff_method="parameter-shift", interface="autograd")
-    qnode_backprop = qml.QNode(circuit, dev, diff_method="backprop", interface="autograd")
 
     # parameter-shift
     t = timeit.repeat("qnode_shift(params)", globals=globals(), number=num, repeat=reps)
@@ -352,15 +363,12 @@ def circuit(params):
     # Gradient timing
     # ===============
 
-    qnode_shift = qml.QNode(circuit, dev, diff_method="parameter-shift", interface="autograd")
-    qnode_backprop = qml.QNode(circuit, dev, diff_method="backprop", interface="autograd")
-
     # parameter-shift
-    t = timeit.repeat("qml.grad(qnode_shift)(params)", globals=globals(), number=num, repeat=reps)
+    t = timeit.repeat("grad_qnode_shift(params)", globals=globals(), number=num, repeat=reps)
     gradient_shift.append([num_params, min(t) / num])
 
     # backprop
-    t = timeit.repeat("qml.grad(qnode_backprop)(params)", globals=globals(), number=num, repeat=reps)
+    t = timeit.repeat("grad_qnode_backprop(params)", globals=globals(), number=num, repeat=reps)
     gradient_backprop.append([num_params, min(t) / num])
 
 gradient_shift = np.array(gradient_shift).T
@@ -399,8 +407,8 @@ def circuit(params):
 # For a better comparison, we can scale the time required for computing the quantum
 # gradients against the time taken for the corresponding forward pass:
 
-gradient_shift[1] /= forward_shift[1, 1:]
-gradient_backprop[1] /= forward_backprop[1, 1:]
+gradient_shift = gradient_shift.at[1].divide(forward_shift[1, 1:])
+gradient_backprop = gradient_backprop.at[1].divide(forward_backprop[1, 1:])
 
 fig, ax = plt.subplots(1, 1, figsize=(6, 4))
 

demonstrations/tutorial_gaussian_transformation.py

@@ -65,10 +65,10 @@
 # ----------------------
 #
 # As before, we import PennyLane, as well as the wrapped version of NumPy provided
-# by PennyLane:
+# by JAX:
 
 import pennylane as qml
-from pennylane import numpy as np
+from jax import numpy as np
 
 ###############################################################################
 # Next, we instantiate a device which will be used to evaluate the circuit.
@@ -102,7 +102,7 @@ def mean_photon_gaussian(mag_alpha, phase_alpha, phi):
 # ------------
 #
 # As in the :ref:`qubit rotation <qubit_rotation>` tutorial, let's now use one
-# of the built-in PennyLane optimizers in order to optimize the quantum circuit
+# of the ``jaxopt`` optimizers in order to optimize the quantum circuit
 # towards the desired output. We want the mean photon number to be exactly one,
 # so we will use a squared-difference cost function:
 
@@ -114,7 +114,7 @@ def cost(params):
 ###############################################################################
 # At the beginning of the optimization, we choose arbitrary small initial parameters:
 
-init_params = np.array([0.015, 0.02, 0.005], requires_grad=True)
+init_params = np.array([0.015, 0.02, 0.005])
 print(cost(init_params))
 
 ###############################################################################
@@ -129,20 +129,23 @@ def cost(params):
 #     We avoided initial parameters which are exactly zero because that
 #     corresponds to a critical point with zero gradient.
 #
-# Now, let's use the :class:`~.pennylane.GradientDescentOptimizer`, and update the circuit
+# Now, let's use the ``GradientDescent`` optimizer, and update the circuit
 # parameters over 100 optimization steps.
 
+import jaxopt
+
 # initialise the optimizer
-opt = qml.GradientDescentOptimizer(stepsize=0.1)
+opt = jaxopt.GradientDescent(cost, stepsize=0.1, acceleration = False)
 
 # set the number of steps
 steps = 20
 # set the initial parameter values
 params = init_params
+opt_state = opt.init_state(params)
 
 for i in range(steps):
     # update the circuit parameters
-    params = opt.step(cost, params)
+    params, opt_state = opt.update(params, opt_state)
 
     print("Cost after step {:5d}: {:8f}".format(i + 1, cost(params)))
 

demonstrations/tutorial_noisy_circuits.metadata.json

@@ -6,7 +6,7 @@
         }
     ],
     "dateOfPublication": "2021-02-22T00:00:00+00:00",
-    "dateOfLastModification": "2021-04-08T00:00:00+00:00",
+    "dateOfLastModification": "2023-09-20T00:00:00+00:00",
     "categories": [
         "Getting Started"
     ],