Update braket static demos; Update torch dependency (#1148)

**Summary:**
* Update `ahs_aquila` and `braket_parallel_gradients` demos to be up to
date with PennyLane 0.37.
* Update `pyproject.toml` to install torch 2.1.2. The old version (1.13)
was quite old. I tried changing the version to 2.3.1 which requires
`typing-extensions>=4.8`, but that caused a dependency conflict with
tensorflow which required `typing-extensions<4.6`

**Relevant references:**

**Possible Drawbacks:**

**Related GitHub Issues:**

Co-authored-by: Astral Cai <[email protected]>

demonstrations/ahs_aquila.py

@@ -254,7 +254,7 @@
 # {'area': {'width': Decimal('0.000075'), 'height': Decimal('0.000076')},
 # 'geometry': {'spacingRadialMin': Decimal('0.000004'),
 # 'spacingVerticalMin': Decimal('0.000004'),
-# 'positionResolution': Decimal('1E-7'),
+# 'positionResolution': Decimal('1E-8'),
 # 'numberSitesMax': 256}}
 #
 # We can see that the atom field has a width of :math:`75 \, \mu m` and a height of :math:`76 \, \mu m`.
@@ -289,6 +289,7 @@
 plt.scatter([x for x, y in coordinates], [y for x, y in coordinates])
 plt.xlabel("μm")
 plt.ylabel("μm")
+plt.show()
 
 ##############################################################################
 # .. figure:: ../_static/demonstration_assets/ahs_aquila/rydberg_blockade_coordinates.png
@@ -347,16 +348,17 @@
 # 'rydbergGlobal': {'rabiFrequencyRange': (Decimal('0.0'),
 # Decimal('15800000.0')),
 # 'rabiFrequencyResolution': Decimal('400.0'),
-# 'rabiFrequencySlewRateMax': Decimal('250000000000000.0'),
+# 'rabiFrequencySlewRateMax': Decimal('400000000000000.0'),
 # 'detuningRange': (Decimal('-125000000.0'), Decimal('125000000.0')),
 # 'detuningResolution': Decimal('0.2'),
-# 'detuningSlewRateMax': Decimal('2500000000000000.0'),
+# 'detuningSlewRateMax': Decimal('6000000000000000.0'),
 # 'phaseRange': (Decimal('-99.0'), Decimal('99.0')),
 # 'phaseResolution': Decimal('5E-7'),
 # 'timeResolution': Decimal('1E-9'),
 # 'timeDeltaMin': Decimal('5E-8'),
 # 'timeMin': Decimal('0.0'),
-# 'timeMax': Decimal('0.000004')}}
+# 'timeMax': Decimal('0.000004')},
+# 'rydbergLocal': None}
 #
 # It is important to note that these quantities are in radians per second rather than Hz where relevant, and
 # are all in SI units. This means that for amplitude and detuning, we will need to convert from angular
@@ -436,7 +438,7 @@ def angular_SI_to_MHz(angular_SI):
 
 
 def gaussian_fn(p, t):
- return p[0] * jnp.exp(-((t-p[1])**2) / (2*p[2]**2))
+ return p[0] * jnp.exp(-((t - p[1]) ** 2) / (2 * p[2] ** 2))
 
 
 # Visualize pulse, time in μs
@@ -453,6 +455,7 @@ def gaussian_fn(p, t):
 plt.ylabel("Amplitude [MHz]")
 
 plt.plot(time, y)
+plt.show()
 
 ##############################################################################
 #
@@ -467,10 +470,7 @@ def gaussian_fn(p, t):
 # We can then define our drive using via :func:`~pennylane.pulse.rydberg_drive`:
 #
 
-global_drive = qml.pulse.rydberg_drive(amplitude=gaussian_fn,
- phase=0,
- detuning=0,
- wires=[0, 1, 2])
+global_drive = qml.pulse.rydberg_drive(amplitude=gaussian_fn, phase=0, detuning=0, wires=[0, 1, 2])
 
 ######################################################################
 # With only amplitude as non-zero, the overall driven Hamiltonian in this case simplifies to:
@@ -606,7 +606,7 @@ def circuit(params):
 start_val = amplitude[0]
 stop_val = amplitude[-1]
 max_val = np.max(amplitude)
-max_rate = np.max([(amplitude[i + 1] - amplitude[i]) / timestep for i in range(len(times)-1)])
+max_rate = np.max([(amplitude[i + 1] - amplitude[i]) / timestep for i in range(len(times) - 1)])
 
 print(f"start value: {start_val:.3} MHz")
 print(f"stop value: {stop_val:.3} MHz")
@@ -636,10 +636,7 @@ def circuit(params):
 #
 
 amp_fn = qml.pulse.rect(gaussian_fn, windows=[0.01, 1.749])
-global_drive = qml.pulse.rydberg_drive(amplitude=amp_fn,
- phase=0,
- detuning=0,
- wires=[0, 1, 2])
+global_drive = qml.pulse.rydberg_drive(amplitude=amp_fn, phase=0, detuning=0, wires=[0, 1, 2])
 
 ######################################################################
 # At this point we could skip directly to defining a ``qnode`` using the ``aquila`` device and running our
@@ -679,6 +676,7 @@ def circuit(params):
 plt.xlabel("μm")
 plt.ylabel("μm")
 plt.legend()
+plt.show()
 
 ##############################################################################
 #
@@ -705,8 +703,7 @@ def circuit(params):
 # values for plotting the function defined in PennyLane for amplitude
 input_times = np.linspace(*ts, 1000)
 input_amplitudes = [
- qml.pulse.rect(gaussian_fn, windows=[0.01, 1.749])(amplitude_params, _t)
- for _t in input_times
+ qml.pulse.rect(gaussian_fn, windows=[0.01, 1.749])(amplitude_params, _t) for _t in input_times
 ]
 
 # plot PL input and hardware setpoints for comparison

demonstrations/braket-parallel-gradients.py

@@ -95,7 +95,6 @@
 Let's load the SV1 simulator in PennyLane with 25 qubits by specifying the device ARN.
 """
 
-
 device_arn = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
 
 ##############################################################################
@@ -193,8 +192,8 @@ def circuit(params):
 #
 # .. code-block:: none
 #
-# Execution time on remote device (seconds): 3.5898206680030853
-# Execution time on local device (seconds): 23.50668462700196
+# Execution time on remote device (seconds): 7.571744918823242
+# Execution time on local device (seconds): 27.32159185409546
 #
 # Nice! These timings highlight the advantage of using the Amazon Braket SV1 device for simulations
 # with large qubit numbers. In general, simulation times scale exponentially with the number of
@@ -235,7 +234,7 @@ def circuit(params):
 #
 # .. code-block:: none
 #
-# Gradient calculation time on remote device (seconds): 20.92005863400118
+# Gradient calculation time on remote device (seconds): 6.4775872230529785
 #
 # Now, the local device:
 #
@@ -260,9 +259,9 @@ def circuit(params):
 #
 # .. code-block:: none
 #
-# Gradient calculation time on local device (seconds): 941.8518133479993
+# Gradient calculation time on local device (seconds): 181.5902488231659
 #
-# Wow, the local device needs around 15 minutes or more! Compare this to less than a minute spent
+# Wow, the local device needs around 3 minutes! Compare this to less around 6 seconds spent
 # calculating the gradient on SV1. This provides a powerful lesson in parallelization.
 #
 # What if we had run on SV1 with ``parallel=False``? It would have taken around 3 minutes—still
@@ -394,37 +393,37 @@ def circuit(params, **kwargs):
 #
 # .. code-block:: none
 #
-# Initial cost: -29.98570234095951
+# Initial cost: -29.98570234095953
 # Completed iteration 1
-# Time to complete iteration: 93.96246099472046 seconds
-# Cost at step 1: -27.154071768632154
+# Time to complete iteration: 11.028863906860352 seconds
+# Cost at step 1: -29.995107300982195
 # Completed iteration 2
-# Time to complete iteration: 84.80994844436646 seconds
-# Cost at step 2: -29.98726230006233
+# Time to complete iteration: 10.132628917694092 seconds
+# Cost at step 2: -29.99981056829394
 # Completed iteration 3
-# Time to complete iteration: 83.13504934310913 seconds
-# Cost at step 3: -29.999163153600062
+# Time to complete iteration: 9.985284090042114 seconds
+# Cost at step 3: -30.00233965117029
 # Completed iteration 4
-# Time to complete iteration: 85.61391234397888 seconds
-# Cost at step 4: -30.002158646044307
+# Time to complete iteration: 10.059210062026978 seconds
+# Cost at step 4: -30.01206799710033
 # Completed iteration 5
-# Time to complete iteration: 86.70688223838806 seconds
-# Cost at step 5: -30.012058444011906
+# Time to complete iteration: 9.966963052749634 seconds
+# Cost at step 5: -30.048670540725286
 # Completed iteration 6
-# Time to complete iteration: 83.26341080665588 seconds
-# Cost at step 6: -30.063709712612443
+# Time to complete iteration: 11.216133832931519 seconds
+# Cost at step 6: -30.13240170461342
 # Completed iteration 7
-# Time to complete iteration: 85.25566911697388 seconds
-# Cost at step 7: -30.32522304705352
+# Time to complete iteration: 10.09446096420288 seconds
+# Cost at step 7: -30.25935979157505
 # Completed iteration 8
-# Time to complete iteration: 83.55433392524719 seconds
-# Cost at step 8: -31.411030331978186
+# Time to complete iteration: 10.020287990570068 seconds
+# Cost at step 8: -30.41499293575487
 # Completed iteration 9
-# Time to complete iteration: 84.08745908737183 seconds
-# Cost at step 9: -33.87153965616938
+# Time to complete iteration: 10.345153093338013 seconds
+# Cost at step 9: -30.587353834845057
 # Completed iteration 10
-# Time to complete iteration: 87.4032838344574 seconds
-# Cost at step 10: -36.05424874438809
+# Time to complete iteration: 10.07306981086731 seconds
+# Cost at step 10: -30.76855713404487
 # Parameters saved to params.npy
 #
 # This example shows us that a 20-qubit QAOA problem can be trained within around 1-2 minutes per
@@ -592,7 +591,7 @@ def cost_function(params, **kwargs):
 # :target: javascript:void(0);
 
 ##############################################################################
-# Great! The loss function gets lower with each iteration, reaching a value of approximately -36.5.
+# Great! The loss function gets lower with each iteration, reaching a value of approximately -37.8.
 #
 # Conclusion
 # ----------------------------------