test_autograd

tests/interfaces/test_autograd.py

@@ -211,6 +211,36 @@ def test_no_batch_transform(self, mocker):
         assert isinstance(res[0], np.ndarray)
         assert res[0].shape == ()
         assert np.allclose(res[0], np.cos(y), atol=0.1)
+
+    @pytest.mark.usefixtures("use_legacy_opmath")
+    def test_no_batch_transform_legacy_opmath(self, mocker):
+        """Test functionality to enable and disable"""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100000)
+
+        H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
+        x = 0.6
+        y = 0.2
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H)
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        spy = mocker.spy(dev, "batch_transform")
+
+        with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"):
+            res = qml.execute([tape], dev, None, device_batch_transform=False)
+
+        spy.assert_not_called()
+
+        res = qml.execute([tape], dev, None, device_batch_transform=True)
+        spy.assert_called()
+
+        assert isinstance(res[0], np.ndarray)
+        assert res[0].shape == ()
+        assert np.allclose(res[0], np.cos(y), atol=0.1)
 
     def test_batch_transform_dynamic_shots(self):
         """Tests that the batch transform considers the number of shots for the execution, not those