fix formatting (V2)

PennyLaneAI · Jul 31, 2024 · 81dcb26 · 81dcb26
1 parent 753da92
commit 81dcb26
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tests/gradients/core/test_vjp.py b/tests/gradients/core/test_vjp.py
@@ -692,6 +692,8 @@ def test_batched_params_probs_jacobian(self):
 
         # Analytically expected Jacobian and VJP
         expected_jac = [-0.5 * np.cos(data) * np.sin(x0 + x1), 0.5 * np.cos(data) * np.sin(x0 + x1)]
-        expected = np.tensordot(expected_jac, dy, axes=[[0, 1], [1, 0]])
-        assert qml.math.shape(vjp) == (1, 2) # num tapes, num trainable tape parameters
-        assert np.allclose(vjp, expected) # Both parameters essentially feed into the same RX rotation
+        expected_vjp = np.tensordot(expected_jac, dy, axes=[[0, 1], [1, 0]])
+        assert qml.math.shape(vjp) == (1, 2)  # num tapes, num trainable tape parameters
+        assert np.allclose(
+            vjp, expected_vjp
+        )  # Both parameters essentially feed into the same RX rotation