add tests for adjoint

Author: refraction-ray

tests/test_circuit.py

@@ -1515,3 +1515,22 @@ def test_get_positional_logical_mapping():
     c.measure_instruction(2)
     c.measure_instruction(0)
     assert c.get_positional_logical_mapping() == {0: 2, 1: 0}
+
+@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
+def test_inverse_jit(backend):
+    K = tc.backend
+    def simple_ansatz(param):
+        c = tc.Circuit(3)
+        for i in range(3):
+            c.cx(i, (i+1)%3)
+            c.rzz(i, (i+1)%3, theta=param[i])
+        c1 = c.inverse()
+        c2 = tc.Circuit(3)
+        c2.x(1)
+        c1.append(c2)
+        return tc.backend.real(c1.expectation_ps(z=[1]))
+
+    v_ansatz = K.jit(K.vvag(simple_ansatz))
+    vs, gs = v_ansatz(K.ones([2, 3], dtype="float32"))
+    assert K.shape_tuple(gs) == (2, 3)
+    np.testing.assert_allclose(K.numpy(vs), -1.*K.ones([2]), atol=1e-5)