fixes to tf tests and runtime improvements

tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Integration tests for using the Autograd interface with shot vectors and with a QNode"""
-# pylint: disable=too-many-arguments
+# pylint: disable=too-many-arguments,redefined-outer-name
 
 import pytest
 
@@ -25,13 +25,17 @@
 shots_and_num_copies = [(((5, 2), 1, 10), 4), ((1, 10, (5, 2)), 4)]
 shots_and_num_copies_hess = [(((5, 1), 10), 2), ((10, (5, 1)), 2)]
 
-SEED_FOR_SPSA = 42
-spsa_kwargs = {"h": 0.05, "num_directions": 20, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
+
+kwargs = {
+    "finite-diff": {"h": 0.05},
+    "parameter-shift": {},
+    "spsa": {"h": 0.05, "num_directions": 20},
+}
 
 qubit_device_and_diff_method = [
-    ["default.qubit.legacy", "finite-diff", {"h": 0.05}],
-    ["default.qubit.legacy", "parameter-shift", {}],
-    ["default.qubit.legacy", "spsa", spsa_kwargs],
+    ["default.qubit.legacy", "finite-diff"],
+    ["default.qubit.legacy", "parameter-shift"],
+    ["default.qubit.legacy", "spsa"],
 ]
 
 TOLS = {
@@ -41,8 +45,16 @@
 }
 
 
+@pytest.fixture
+def gradient_kwargs(request):
+    diff_method = request.node.funcargs["diff_method"]
+    return kwargs[diff_method] | (
+        {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {}
+    )
+
+
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)
-@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method)
 class TestReturnWithShotVectors:
     """Class to test the shape of the Jacobian/Hessian with different return types and shot vectors."""
 
@@ -369,7 +381,7 @@ def cost(a):
 
 @pytest.mark.slow
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess)
-@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method)
 class TestReturnShotVectorHessian:
     """Class to test the shape of the Hessian with different return types and shot vectors."""
 
@@ -553,12 +565,11 @@ def cost2(x):
         assert hess.shape == (num_copies, 3, 2, 2)
 
 
-shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)]
+shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)]
 
 
-@pytest.mark.skip("failing in CI for inscrutable reasons, passes locally")
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)
-@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method)
 class TestReturnShotVectorIntegration:
     """Tests for the integration of shots with the autograd interface."""
 

tests/interfaces/legacy_devices_integration/test_tensorflow_autograph_qnode_shot_vector_legacy.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Integration tests for using the TF interface with shot vectors and with a QNode"""
-# pylint: disable=too-many-arguments,too-few-public-methods
+# pylint: disable=too-many-arguments,too-few-public-methods,redefined-outer-name
 import pytest
 
 import pennylane as qml
@@ -26,11 +26,17 @@
 shots_and_num_copies = [((1, (5, 2), 10), 4)]
 shots_and_num_copies_hess = [((10, (5, 1)), 2)]
 
-spsa_kwargs = {"h": 10e-2, "num_directions": 30, "sampler_rng": np.random.default_rng(42)}
+
+kwargs = {
+    "finite-diff": {"h": 10e-2},
+    "parameter-shift": {},
+    "spsa": {"h": 10e-2, "num_directions": 30},
+}
+
 qubit_device_and_diff_method = [
-    ["default.qubit.legacy", "finite-diff", {"h": 10e-2}],
-    ["default.qubit.legacy", "parameter-shift", {}],
-    ["default.qubit.legacy", "spsa", spsa_kwargs],
+    ["default.qubit.legacy", "finite-diff"],
+    ["default.qubit.legacy", "parameter-shift"],
+    ["default.qubit.legacy", "spsa"],
 ]
 
 TOLS = {
@@ -40,8 +46,16 @@
 }
 
 
+@pytest.fixture
+def gradient_kwargs(request):
+    diff_method = request.node.funcargs["diff_method"]
+    return kwargs[diff_method] | (
+        {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {}
+    )
+
+
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)
-@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method)
 @pytest.mark.parametrize(
     "decorator,interface",
     [(tf.function, "tf"), (lambda x: x, "tf-autograph")],
@@ -337,7 +351,7 @@ def circuit(a):
 
 @pytest.mark.slow
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess)
-@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method)
 @pytest.mark.parametrize(
     "decorator,interface",
     [(tf.function, "tf"), (lambda x: x, "tf-autograph")],
@@ -388,7 +402,7 @@ def circuit(x, y):
 
 
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)
-@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method)
 @pytest.mark.parametrize(
     "decorator,interface",
     [(tf.function, "tf"), (lambda x: x, "tf-autograph")],

tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_shot_vector_legacy.py

@@ -453,7 +453,7 @@ def circuit(x):
         assert hess.shape == (num_copies, 3, 2, 2)
 
 
-shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)]
+shots_and_num_copies = [((20000, 18000, 16000), 3), ((20000, (18000, 2)), 3)]
 
 
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)

tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Integration tests for using the TF interface with shot vectors and with a QNode"""
-# pylint: disable=too-many-arguments,too-few-public-methods,unexpected-keyword-arg
+# pylint: disable=too-many-arguments,too-few-public-methods,unexpected-keyword-arg,redefined-outer-name
 import pytest
 
 import pennylane as qml
@@ -27,14 +27,16 @@
 shots_and_num_copies = [((1, (5, 2), 10), 4)]
 shots_and_num_copies_hess = [((10, (5, 1)), 2)]
 
+kwargs = {
+    "finite-diff": {"h": 10e-2},
+    "parameter-shift": {},
+    "spsa": {"h": 10e-2, "num_directions": 20},
+}
+
 qubit_device_and_diff_method = [
-    [DefaultQubit(), "finite-diff", {"h": 10e-2}],
-    [DefaultQubit(), "parameter-shift", {}],
-    [
-        DefaultQubit(),
-        "spsa",
-        {"h": 10e-2, "num_directions": 20, "sampler_rng": np.random.default_rng(42)},
-    ],
+    [DefaultQubit(), "finite-diff"],
+    [DefaultQubit(), "parameter-shift"],
+    [DefaultQubit(), "spsa"],
 ]
 
 TOLS = {
@@ -44,8 +46,16 @@
 }
 
 
+@pytest.fixture
+def gradient_kwargs(request):
+    diff_method = request.node.funcargs["diff_method"]
+    return kwargs[diff_method] | (
+        {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {}
+    )
+
+
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)
-@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method)
 @pytest.mark.parametrize(
     "decorator,interface",
     [(tf.function, "tf"), (lambda x: x, "tf-autograph")],
@@ -329,7 +339,7 @@ def circuit(a, **_):
 
 @pytest.mark.slow
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess)
-@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method)
 @pytest.mark.parametrize(
     "decorator,interface",
     [(tf.function, "tf"), (lambda x: x, "tf-autograph")],
@@ -374,11 +384,11 @@ def circuit(x, y, **_):
             assert h.shape == (2, num_copies)
 
 
-shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)]
+shots_and_num_copies = [((30000, 28000, 26000), 3), ((20000, (18000, 2)), 3)]
 
 
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)
-@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method)
+@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method)
 @pytest.mark.parametrize(
     "decorator,interface",
     [(tf.function, "tf"), (lambda x: x, "tf-autograph")],

tests/interfaces/test_tensorflow_qnode_shot_vector.py

@@ -438,7 +438,7 @@ def circuit(x):
         assert hess.shape == (num_copies, 3, 2, 2)
 
 
-shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)]
+shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)]
 
 
 @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies)