def ideal_experiment(): qml.SX(wires=0) return qml.state()
def my_circuit(): qml.adjoint(qml.Barrier)(wires=0) return qml.state()
def my_circuit(): adjoint(my_op)() qml.Hadamard(wires=0) adjoint(adjoint(my_op))() return qml.state()
def circuit(features): template(features=features, wires=[0, 1, 2]) fn(template, features=features, wires=[0, 1, 2]) return qml.state()
def circuit(weights): template(weight=weights, wires=[0, 1, 2]) fn(template, weight=weights, wires=[0, 1, 2]) return qml.state()
def circuit_template(features): qml.templates.BasisEmbedding(features, wires=range(3)) return qml.state()
def qfunc1(): qml.RX(2, wires=0) qml.RY(-3, wires=1) qml.QFT(wires=[0, 1, 2]) qml.SWAP(wires=[1, 2]) return qml.state()
def circuit(): qml.Hadamard(wires=0) qml.RZ(jnp.pi / 4, wires=0) return qml.state()
def circuit(a): qml.RY(a, wires=0) return qml.state()
def circuit(): qml.RY(0.5, wires=0) return qml.state()
def get_state(): circuit() return qml.state()
def cost(a, b, device): with JAXInterface.apply(JacobianTape()) as tape: qml.RY(a, wires=0) qml.RX(b, wires=0) qml.state() return tape.execute(device)
_, ax = tape_mpl(tape, decimals=2) num_wires = len(op.wires) assert ax.texts[num_wires].get_text() == op.label(decimals=2) plt.close() measure_data = [ ([qml.expval(qml.PauliX(0))], [0]), ([qml.probs(wires=(0, 1, 2))], [0, 1, 2]), ([ qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0) @ qml.PauliY(1)), qml.state() ], [0, 1]), ([qml.expval(qml.NumberOperator(wires=0))], [0]), ] class TestMeasurements: """Tests measurements are drawn correctly""" @pytest.mark.parametrize("measurements, wires", measure_data) def test_measurements(self, measurements, wires): """Tests a variety of measurements draw measurement boxes on the correct wires.""" with QuantumTape() as tape: for m in measurements: qml.apply(m)
def circuit(weight): qml.GateFabric(weight, wires, init_state=init_state, include_pi=True) return qml.state()
def circuit(ansatz, params): qml.RX(np.pi / 4.0, wires=2) ansatz(params) return qml.state()
def output(input): qml.BasisState(input, wires=range(wires)) op(*p, wires=range(wires)).inv() return qml.state()
def circuit(): qml.AmplitudeEmbedding(np.array([0, 1]), wires=0) return qml.state()
def output(input): qml.BasisState(input, wires=range(wires)) qml.QubitUnitary(random_unitary, wires=range(2)).inv() return qml.state()
def qfunc(): qml.PauliX(wires=1) qml.CNOT(wires=[0, 1]) qml.SWAP(wires=[0, 1]) return qml.state()
def my_circuit(): my_op() invisible(my_op)() qml.Hadamard(wires=0) invisible(invisible(my_op))() return qml.state()
def qfunc2(): qml.RX(2, wires=0) qml.RY(-3, wires=2) qml.QFT(wires=[0, 2, 1]) return qml.state()
def circuit(): qml.QubitDensityMatrix(initialize_state, wires=[0, 1]) return qml.state()
def circuit(t): template(H, t, 1) fn(template, H, t, 1) return qml.state()
class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" def test_execution_with_interface(self, dev_name, diff_method, mode): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") dev = qml.device(dev_name, wires=1) @qnode(dev, interface="jax", diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) a = np.array(0.1, requires_grad=True) circuit(a) assert circuit.interface == "jax" # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] # gradients should work grad = jax.grad(circuit)(a) assert isinstance(grad, jnp.DeviceArray) assert grad.shape == tuple() def test_jacobian(self, dev_name, diff_method, mode, mocker, tol): """Test jacobian calculation""" if diff_method != "backprop": pytest.skip("JAX interface does not support vector-valued QNodes") if diff_method == "parameter-shift": spy = mocker.spy(qml.gradients.param_shift, "transform_fn") elif diff_method == "finite-diff": spy = mocker.spy(qml.gradients.finite_diff, "transform_fn") a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) dev = qml.device(dev_name, wires=2) @qnode(dev, diff_method=diff_method, interface="jax", mode=mode) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] res = circuit(a, b) assert circuit.qtape.trainable_params == [0, 1] assert res.shape == (2,) expected = [np.cos(a), -np.cos(a) * np.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) res = jax.jacobian(circuit, argnums=[0, 1])(a, b) expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T assert np.allclose(res, expected, atol=tol, rtol=0) if diff_method in ("parameter-shift", "finite-diff"): spy.assert_called() def test_jacobian_no_evaluate(self, dev_name, diff_method, mode, mocker, tol): """Test jacobian calculation when no prior circuit evaluation has been performed""" if diff_method != "backprop": pytest.skip("JAX interface does not support vector-valued QNodes") if diff_method == "parameter-shift": spy = mocker.spy(qml.gradients.param_shift, "transform_fn") elif diff_method == "finite-diff": spy = mocker.spy(qml.gradients.finite_diff, "transform_fn") a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) dev = qml.device(dev_name, wires=2) @qnode(dev, diff_method=diff_method, interface="jax", mode=mode) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] jac_fn = jax.jacobian(circuit, argnums=[0, 1]) res = jac_fn(a, b) expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T assert np.allclose(res, expected, atol=tol, rtol=0) if diff_method in ("parameter-shift", "finite-diff"): spy.assert_called() # call the Jacobian with new parameters a = np.array(0.6, requires_grad=True) b = np.array(0.832, requires_grad=True) res = jac_fn(a, b) expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T assert np.allclose(res, expected, atol=tol, rtol=0) def test_jacobian_options(self, dev_name, diff_method, mode, mocker, tol): """Test setting jacobian options""" if diff_method == "backprop": pytest.skip("Test does not support backprop") spy = mocker.spy(qml.gradients.finite_diff, "transform_fn") a = np.array([0.1, 0.2], requires_grad=True) dev = qml.device("default.qubit", wires=1) @qnode(dev, interface="jax", h=1e-8, order=2) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) jax.jacobian(circuit)(a) for args in spy.call_args_list: assert args[1]["order"] == 2 assert args[1]["h"] == 1e-8 def test_changing_trainability(self, dev_name, diff_method, mode, mocker, tol): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": pytest.skip("Test only supports parameter-shift") a = jnp.array(0.1) b = jnp.array(0.2) dev = qml.device("default.qubit", wires=2) @qnode(dev, interface="jax", diff_method="parameter-shift") def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliY(1)])) grad_fn = jax.grad(circuit, argnums=[0, 1]) spy = mocker.spy(qml.gradients.param_shift, "transform_fn") res = grad_fn(a, b) # the tape has reported both arguments as trainable assert circuit.qtape.trainable_params == [0, 1] expected = [-np.sin(a) + np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=tol, rtol=0) # The parameter-shift rule has been called for each argument assert len(spy.spy_return[0]) == 4 # make the second QNode argument a constant grad_fn = jax.grad(circuit, argnums=0) res = grad_fn(a, b) # the tape has reported only the first argument as trainable assert circuit.qtape.trainable_params == [0] expected = [-np.sin(a) + np.sin(a) * np.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) # The parameter-shift rule has been called only once assert len(spy.spy_return[0]) == 2 # trainability also updates on evaluation a = np.array(0.54, requires_grad=False) b = np.array(0.8, requires_grad=True) circuit(a, b) assert circuit.qtape.trainable_params == [1] def test_classical_processing(self, dev_name, diff_method, mode, tol): """Test classical processing within the quantum tape""" a = jnp.array(0.1) b = jnp.array(0.2) c = jnp.array(0.3) dev = qml.device(dev_name, wires=1) @qnode(dev, diff_method=diff_method, interface="jax", mode=mode) def circuit(a, b, c): qml.RY(a * c, wires=0) qml.RZ(b, wires=0) qml.RX(c + c ** 2 + jnp.sin(a), wires=0) return qml.expval(qml.PauliZ(0)) res = jax.grad(circuit, argnums=[0, 2])(a, b, c) if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [0, 2] assert len(res) == 2 def test_matrix_parameter(self, dev_name, diff_method, mode, tol): """Test that the jax interface works correctly with a matrix parameter""" U = jnp.array([[0, 1], [1, 0]]) a = jnp.array(0.1) dev = qml.device(dev_name, wires=2) @qnode(dev, diff_method=diff_method, interface="jax", mode=mode) def circuit(U, a): qml.QubitUnitary(U, wires=0) qml.RY(a, wires=0) return qml.expval(qml.PauliZ(0)) res = jax.grad(circuit, argnums=1)(U, a) assert np.allclose(res, np.sin(a), atol=tol, rtol=0) if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] def test_differentiable_expand(self, dev_name, diff_method, mode, tol): """Test that operation and nested tape expansion is differentiable""" class U3(qml.U3): def expand(self): theta, phi, lam = self.data wires = self.wires with JacobianTape() as tape: qml.Rot(lam, theta, -lam, wires=wires) qml.PhaseShift(phi + lam, wires=wires) return tape dev = qml.device(dev_name, wires=1) a = jnp.array(0.1) p = jnp.array([0.1, 0.2, 0.3]) @qnode(dev, diff_method=diff_method, interface="jax", mode=mode) def circuit(a, p): qml.RX(a, wires=0) U3(p[0], p[1], p[2], wires=0) return qml.expval(qml.PauliX(0)) res = circuit(a, p) expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) ) assert np.allclose(res, expected, atol=tol, rtol=0) res = jax.grad(circuit, argnums=1)(a, p) expected = np.array( [ np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), np.cos(p[1]) * np.cos(p[2]) * np.sin(a) - np.sin(p[1]) * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), np.sin(a) * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), ] ) assert np.allclose(res, expected, atol=tol, rtol=0) def test_multiple_outputs_raises(self, dev_name, diff_method, mode, tol): """Test executing a QNode that has multiple outputs raises an error.""" dev = qml.device(dev_name, wires=2) if diff_method == "backprop": pytest.skip("Test is not applicable for backprop") @qml.qnode(dev, interface="jax", diff_method=diff_method, mode=mode) def my_circuit(param): qml.RX(param, wires=0) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) with pytest.raises( ValueError, match="JAX interface currently only supports quantum nodes with a single return type", ): my_circuit(1) @pytest.mark.parametrize("ret", [qml.probs(wires=0), qml.state()]) def test_not_expval_or_var_raises(self, dev_name, diff_method, mode, ret, tol): """Test executing a QNode that has a return type other than expval or var raises an error.""" dev = qml.device(dev_name, wires=2) if diff_method == "backprop": pytest.skip("Test is not applicable for backprop") if diff_method == "adjoint": pytest.skip("Adjoint does not support states") @qml.qnode(dev, interface="jax", diff_method=diff_method, mode=mode) def my_circuit(param): qml.RX(param, wires=0) qml.CNOT(wires=[0, 1]) return qml.apply(ret) with pytest.raises( ValueError, match="Only Variance and Expectation returns are supported for the JAX interface", ): my_circuit(1)
def circuit(weights): template(weight=weights, wires1=[0, 1], wires2=[2, 3]) fn(template, weight=weights, wires1=[0, 1], wires2=[2, 3]) return qml.state()
def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.state()
def my_circuit(): adjoint(adjoint(qml.RX))(np.pi / 4.0, wires=0) return qml.state()
def circuit(): qml.PauliX(wires=0) return qml.state()
def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) qml.RY(p2[0] * p2[1], wires=1) qml.RX(kwargs["p3"], wires=0) qml.CNOT(wires=[0, 1]) return qml.state()
def circuit_decomposed(features): # need to cast to complex tensor, which is implicitly done in the template qml.QubitStateVector(qml.math.cast(features, np.complex128), wires=range(3)) return qml.state()