diff --git a/demonstrations/adjoint_diff_benchmarking.py b/demonstrations/adjoint_diff_benchmarking.py
index 989bf92322..e1369bc77e 100644
--- a/demonstrations/adjoint_diff_benchmarking.py
+++ b/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)
diff --git a/demonstrations/tutorial_adjoint_diff.py b/demonstrations/tutorial_adjoint_diff.py
index 94bbc84b26..58e89dc41e 100644
--- a/demonstrations/tutorial_adjoint_diff.py
+++ b/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.
 #
diff --git a/demonstrations/tutorial_backprop.py b/demonstrations/tutorial_backprop.py
index a7c612b313..9a23e09837 100644
--- a/demonstrations/tutorial_backprop.py
+++ b/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,7 +127,7 @@ 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
 
@@ -129,12 +135,11 @@ def parameter_shift(qnode, 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))
 
diff --git a/demonstrations/tutorial_gaussian_transformation.py b/demonstrations/tutorial_gaussian_transformation.py
index 70e531f0e1..39176862f9 100644
--- a/demonstrations/tutorial_gaussian_transformation.py
+++ b/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)))
 
diff --git a/demonstrations/tutorial_noisy_circuits.metadata.json b/demonstrations/tutorial_noisy_circuits.metadata.json
index fe6368f6e5..4562549534 100644
--- a/demonstrations/tutorial_noisy_circuits.metadata.json
+++ b/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"
     ],
diff --git a/demonstrations/tutorial_noisy_circuits.py b/demonstrations/tutorial_noisy_circuits.py
index 9f62f443e2..ab4d6a11da 100644
--- a/demonstrations/tutorial_noisy_circuits.py
+++ b/demonstrations/tutorial_noisy_circuits.py
@@ -60,10 +60,14 @@
 # return the expectation value of :math:`Z_0\otimes Z_1`:
 #
 import pennylane as qml
-from pennylane import numpy as np
+from jax import numpy as np
+import jax
+import jaxopt
 
-dev = qml.device('default.mixed', wires=2)
+jax.config.update("jax_platform_name", "cpu")
+jax.config.update('jax_enable_x64', True)
 
+dev = qml.device('default.mixed', wires=2)
 
 @qml.qnode(dev)
 def circuit():
@@ -262,17 +266,14 @@ def cost(x, target):
 # All that remains is to optimize the parameter. We use a straightforward gradient descent
 # method.
 
-
-opt = qml.GradientDescentOptimizer(stepsize=10)
 steps = 35
-x = np.tensor(0.01, requires_grad=True)
 
-for i in range(steps):
-    (x, ev), cost_val = opt.step_and_cost(cost, x, ev)
-    if i % 5 == 0 or i == steps - 1:
-        print(f"Step: {i}    Cost: {cost_val}")
+gd = jaxopt.GradientDescent(cost, maxiter=steps)
+
+x = np.array(0.01)
+res = gd.run(x, ev)
 
-print(f"QNode output after optimization = {damping_circuit(x):.4f}")
+print(f"QNode output after optimization = {damping_circuit(res.params):.4f}")
 print(f"Experimental expectation value = {ev}")
 print(f"Optimized noise parameter p = {sigmoid(x.take(0)):.4f}")
 
diff --git a/demonstrations/tutorial_phase_kickback.py b/demonstrations/tutorial_phase_kickback.py
index dced30a476..db07460891 100644
--- a/demonstrations/tutorial_phase_kickback.py
+++ b/demonstrations/tutorial_phase_kickback.py
@@ -41,12 +41,12 @@
 # Setting up PennyLane
 # --------------------
 #
-# First, let’s import the necessary PennyLane libraries and create a device to run our quantum
+# First, let’s import the necessary libraries and create a device to run our quantum
 # circuits. Here we will work with 5 qubits, we will use qubit [0] as the control ancilla qubit, and qubits [1,2,3,4] will be our target qubits where we will encode :math:`|\psi\rangle`.
 #
 
 import pennylane as qml
-from pennylane import numpy as np
+import numpy as np
 
 num_wires = 5
 dev = qml.device("default.qubit", wires=num_wires, shots=1)
diff --git a/demonstrations/tutorial_qft_arithmetics.py b/demonstrations/tutorial_qft_arithmetics.py
index 05d97e4f46..9bddca9997 100644
--- a/demonstrations/tutorial_qft_arithmetics.py
+++ b/demonstrations/tutorial_qft_arithmetics.py
@@ -172,7 +172,7 @@ def basis_embedding_circuit(m):
 #
 
 import pennylane as qml
-from pennylane import numpy as np
+import numpy as np
 
 n_wires = 4
 dev = qml.device("default.qubit", wires=n_wires, shots=1)
diff --git a/demonstrations/tutorial_qubit_rotation.py b/demonstrations/tutorial_qubit_rotation.py
index 1ff569f849..feba97ddfe 100644
--- a/demonstrations/tutorial_qubit_rotation.py
+++ b/demonstrations/tutorial_qubit_rotation.py
@@ -86,26 +86,11 @@
 # -----------------------------
 #
 # The first thing we need to do is import PennyLane, as well as the wrapped version
-# of NumPy provided by PennyLane.
+# of NumPy provided by Jax.
 
 import pennylane as qml
-from pennylane import numpy as np
-
-
-##############################################################################
-# .. important::
-#
-#     When constructing a hybrid quantum/classical computational model with PennyLane,
-#     it is important to **always import NumPy from PennyLane**, not the standard NumPy!
-#
-#     By importing the wrapped version of NumPy provided by PennyLane, you can combine
-#     the power of NumPy with PennyLane:
-#
-#     * continue to use the classical NumPy functions and arrays you know and love
-#     * combine quantum functions (evaluated on quantum hardware/simulators) and
-#       classical functions (provided by NumPy)
-#     * allow PennyLane to automatically calculate gradients of both classical and
-#       quantum functions
+from jax import numpy as np
+import jax
 
 
 ##############################################################################
@@ -120,9 +105,7 @@
 #     Any computational object that can apply quantum operations and return a measurement value
 #     is called a quantum **device**.
 #
-#     In PennyLane, a device could be a hardware device (such as the IBM QX4, via the
-#     PennyLane-PQ plugin), or a software simulator (such as Strawberry Fields, via the
-#     PennyLane-SF plugin).
+#     In PennyLane, a device could be a hardware device (take a look at our `plugins <https://pennylane.ai/plugins/#plugins>`_), or a software simulator (such as our high performance simulator `PennyLane-Lightning <https://docs.pennylane.ai/projects/lightning/en/stable/>`_).
 #
 # .. tip::
 #
@@ -214,7 +197,7 @@ def circuit(params):
 # **directly above** the function definition:
 
 
-@qml.qnode(dev1, interface="autograd")
+@qml.qnode(dev1)
 def circuit(params):
     qml.RX(params[0], wires=0)
     qml.RY(params[1], wires=0)
@@ -227,7 +210,8 @@ def circuit(params):
 #
 # To evaluate, we simply call the function with some appropriate numerical inputs:
 
-print(circuit([0.54, 0.12]))
+params = np.array([0.54, 0.12])
+print(circuit(params))
 
 ##############################################################################
 # Calculating quantum gradients
@@ -241,21 +225,21 @@ def circuit(params):
 # methods (such as the method of finite differences). Both of these are done
 # automatically.
 #
-# We can differentiate by using the built-in :func:`~.pennylane.grad` function.
+# We can differentiate by using the `jax.grad` function.
 # This returns another function, representing the gradient (i.e., the vector of
 # partial derivatives) of ``circuit``. The gradient can be evaluated in the same
 # way as the original function:
 
-dcircuit = qml.grad(circuit, argnum=0)
+dcircuit = jax.grad(circuit, argnums=0)
 
 ##############################################################################
-# The function :func:`~.pennylane.grad` itself **returns a function**, representing
-# the derivative of the QNode with respect to the argument specified in ``argnum``.
+# The function `jax.grad` itself **returns a function**, representing
+# the derivative of the QNode with respect to the argument specified in ``argnums``.
 # In this case, the function ``circuit`` takes one argument (``params``), so we
-# specify ``argnum=0``. Because the argument has two elements, the returned gradient
+# specify ``argnums=0``. Because the argument has two elements, the returned gradient
 # is two-dimensional. We can then evaluate this gradient function at any point in the parameter space.
 
-print(dcircuit([0.54, 0.12]))
+print(dcircuit(params))
 
 ################################################################################
 # **A note on arguments**
@@ -266,7 +250,7 @@ def circuit(params):
 # two positional arguments, instead of one array argument:
 
 
-@qml.qnode(dev1, interface="autograd")
+@qml.qnode(dev1)
 def circuit2(phi1, phi2):
     qml.RX(phi1, wires=0)
     qml.RY(phi2, wires=0)
@@ -274,14 +258,17 @@ def circuit2(phi1, phi2):
 
 
 ################################################################################
-# When we calculate the gradient for such a function, the usage of ``argnum``
-# will be slightly different. In this case, ``argnum=0`` will return the gradient
-# with respect to only the first parameter (``phi1``), and ``argnum=1`` will give
+# When we calculate the gradient for such a function, the usage of ``argnums``
+# will be slightly different. In this case, ``argnums=0`` will return the gradient
+# with respect to only the first parameter (``phi1``), and ``argnums=1`` will give
 # the gradient for ``phi2``. To get the gradient with respect to both parameters,
-# we can use ``argnum=[0,1]``:
+# we can use ``argnums=[0,1]``:
 
-dcircuit = qml.grad(circuit2, argnum=[0, 1])
-print(dcircuit(0.54, 0.12))
+phi1 = np.array([0.54])
+phi2 = np.array([0.12])
+
+dcircuit = jax.grad(circuit2, argnums=[0, 1])
+print(dcircuit(phi1, phi2))
 
 ################################################################################
 # Keyword arguments may also be used in your custom quantum function. PennyLane
@@ -334,26 +321,29 @@ def cost(x):
 ################################################################################
 # To begin our optimization, let's choose small initial values of :math:`\phi_1` and :math:`\phi_2`:
 
-init_params = np.array([0.011, 0.012], requires_grad=True)
+init_params = np.array([0.011, 0.012])
 print(cost(init_params))
 
 ################################################################################
 # We can see that, for these initial parameter values, the cost function is close to :math:`1`.
 #
-# Finally, we use an optimizer to update the circuit parameters for 100 steps. We can use the built-in
-# :class:`~.pennylane.GradientDescentOptimizer` class:
+# Finally, we use an optimizer to update the circuit parameters for 100 steps. We can use the
+# gradient descent optimizer:
+
+import jaxopt
 
 # initialise the optimizer
-opt = qml.GradientDescentOptimizer(stepsize=0.4)
+opt = jaxopt.GradientDescent(cost, stepsize=0.4, acceleration = False)
 
 # set the number of steps
 steps = 100
 # 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)
 
     if (i + 1) % 5 == 0:
         print("Cost after step {:5d}: {: .7f}".format(i + 1, cost(params)))
diff --git a/demonstrations/tutorial_vqe.py b/demonstrations/tutorial_vqe.py
index b120e30015..cb30eb0b23 100644
--- a/demonstrations/tutorial_vqe.py
+++ b/demonstrations/tutorial_vqe.py
@@ -47,7 +47,10 @@
 and a one-dimensional array with the corresponding nuclear coordinates
 in `atomic units <https://en.wikipedia.org/wiki/Hartree_atomic_units>`_.
 """
-from pennylane import numpy as np
+from jax import numpy as np
+import jax
+jax.config.update("jax_platform_name", "cpu")
+jax.config.update('jax_enable_x64', True)
 
 symbols = ["H", "H"]
 coordinates = np.array([0.0, 0.0, -0.6614, 0.0, 0.0, 0.6614])
@@ -146,45 +149,46 @@
 ##############################################################################
 # The ``hf`` array is used by the :class:`~.pennylane.BasisState` operation to initialize
 # the qubit register. Then, we just act with the :class:`~.pennylane.DoubleExcitation` operation
-# on the four qubits.
-
+# on the four qubits. The next step is to compute the expectation value
+# of the molecular Hamiltonian in the trial state prepared by the circuit.
+# We do this using the :func:`~.expval` function. The decorator syntax allows us to
+# run the cost function as an executable QNode with the gate parameter :math:`\theta`:
 
+@qml.qnode(dev)
 def circuit(param, wires):
     qml.BasisState(hf, wires=wires)
     qml.DoubleExcitation(param, wires=[0, 1, 2, 3])
+    return qml.expval(H)
 
 
 ##############################################################################
-# The next step is to define the cost function to compute the expectation value
-# of the molecular Hamiltonian in the trial state prepared by the circuit.
-# We do this using the :func:`~.expval` function. The decorator syntax allows us to
-# run the cost function as an executable QNode with the gate parameter :math:`\theta`:
+# We can now define our error function simply as the expected value calculated above:
 
 
-@qml.qnode(dev, interface="autograd")
 def cost_fn(param):
-    circuit(param, wires=range(qubits))
-    return qml.expval(H)
-
+    return circuit(param, wires=range(qubits))
 
 ##############################################################################
 # Now we proceed to minimize the cost function to find the ground state of
 # the :math:`\mathrm{H}_2` molecule. To start, we need to define the classical optimizer.
-# PennyLane offers many different built-in
-# `optimizers <https://docs.pennylane.ai/en/stable/introduction/interfaces.html?highlight=optimizers#optimizers>`_.
+# The library ``optax`` offers different `optimizers <https://optax.readthedocs.io/en/latest/api.html>`__.
 # Here we use a basic gradient-descent optimizer.
+# We carry out the optimization over a maximum of 100 steps aiming to reach a
+# convergence tolerance of :math:`10^{-6}` for the value of the cost function.
+
 
-opt = qml.GradientDescentOptimizer(stepsize=0.4)
+import optax
+
+max_iterations = 100
+conv_tol = 1e-06
+
+opt = optax.sgd(learning_rate=0.4)
 
 ##############################################################################
 # We initialize the circuit parameter :math:`\theta` to zero, meaning that we start
 # from the Hartree-Fock state.
 
-theta = np.array(0.0, requires_grad=True)
-
-##############################################################################
-# We carry out the optimization over a maximum of 100 steps aiming to reach a
-# convergence tolerance of :math:`10^{-6}` for the value of the cost function.
+theta = np.array(0.)
 
 # store the values of the cost function
 energy = [cost_fn(theta)]
@@ -192,16 +196,18 @@ def cost_fn(param):
 # store the values of the circuit parameter
 angle = [theta]
 
-max_iterations = 100
-conv_tol = 1e-06
+opt_state = opt.init(theta)
 
 for n in range(max_iterations):
-    theta, prev_energy = opt.step_and_cost(cost_fn, theta)
 
-    energy.append(cost_fn(theta))
+    gradient = jax.grad(cost_fn)(theta)
+    updates, opt_state = opt.update(gradient, opt_state)
+    theta = optax.apply_updates(theta, updates)
+    
     angle.append(theta)
+    energy.append(cost_fn(theta))
 
-    conv = np.abs(energy[-1] - prev_energy)
+    conv = np.abs(energy[-1] - energy[-2])
 
     if n % 2 == 0:
         print(f"Step = {n},  Energy = {energy[-1]:.8f} Ha")