def test_tf_not_trainable_only(diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Tensorflow with ``trainable_only=False`` .""" tf = pytest.importorskip("tensorflow") dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit_0, dev, interface="tf", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=0, trainable_only=False)(tf.Variable(a)) assert np.allclose(jac, expected_jac_not_trainable_only)
def test_jax_not_trainable_only(diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and JAX with ``trainable_only=False`` .""" # Do not need the package but skip if JAX device not available pytest.importorskip("jax") dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit_0, dev, interface="jax", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=0, trainable_only=False)(a) assert np.allclose(jac, expected_jac_not_trainable_only)
def test_tf_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and TensorFlow.""" tf = pytest.importorskip("tensorflow") args = tuple((tf.Variable(arg, dtype=tf.double) for arg in args)) dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="tf", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = expected_jac[argnum] assert np.allclose(jac, expected_jac)
def test_jax_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and JAX.""" # Do not need the package but skip if JAX device not available pytest.importorskip("jax") dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = expected_jac[argnum] assert np.allclose(jac, expected_jac)
def test_torch_without_argnum(circuit, args, expected_jac, diff_method): r"""Test ``classical_jacobian`` with ``argnum=None`` and Torch.""" torch = pytest.importorskip("torch") args = tuple((torch.tensor(arg, requires_grad=True) for arg in args)) dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method) jac = classical_jacobian(qnode)(*args) assert len(jac) == len(expected_jac) for _jac, _expected_jac in zip(jac, expected_jac): assert np.allclose(_jac, _expected_jac) # also test with an untrainable argument args[0].requires_grad = False jac = classical_jacobian(qnode)(*args) assert len(jac) == len(expected_jac) - 1 for _jac, _expected_jac in zip(jac, expected_jac[1:]): assert np.allclose(_jac, _expected_jac)
def test_jax_without_argnum(circuit, args, expected_jac, diff_method): r"""Test ``classical_jacobian`` with ``argnum=None`` and JAX.""" # Do not need the package but skip if JAX device not available pytest.importorskip("jax") # JAX behaviour: argnum=None yields only the Jacobian with respect to the first arg. expected_jac = expected_jac[0] dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method) jac = classical_jacobian(qnode)(*args) assert np.allclose(jac, expected_jac)
def test_autograd_not_trainable_only(diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Autograd with ``trainable_only=False`` .""" dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit_0, dev, interface="autograd", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=0, trainable_only=False)(a) assert np.allclose(jac, expected_jac_not_trainable_only)
def test_torch_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Torch.""" torch = pytest.importorskip("torch") args = tuple((torch.tensor(arg) for arg in args)) dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = expected_jac[argnum] assert np.allclose(jac, expected_jac)
def test_torch_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and Torch.""" torch = pytest.importorskip("torch") args = tuple((torch.tensor(arg) for arg in args)) dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = (expected_jac[argnum[0]], ) assert len(jac) == 1 assert np.allclose(jac[0], expected_jac[0])
def test_jax_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and JAX.""" # Do not need the package but skip if JAX device not available pytest.importorskip("jax") dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = (expected_jac[argnum[0]], ) assert len(jac) == 1 assert np.allclose(jac[0], expected_jac[0])
def test_autograd_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Autograd.""" dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = expected_jac[argnum] assert np.allclose(jac, expected_jac)
def test_tf_without_argnum(circuit, args, expected_jac, diff_method): r"""Test ``classical_jacobian`` with ``argnum=None`` and Tensorflow.""" tf = pytest.importorskip("tensorflow") args = tuple((tf.Variable(arg, dtype=tf.double) for arg in args)) dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="tf", diff_method=diff_method) jac = classical_jacobian(qnode)(*args) assert len(jac) == len(expected_jac) for _jac, _expected_jac in zip(jac, expected_jac): assert np.allclose(_jac, _expected_jac)
def test_torch_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and Torch.""" torch = pytest.importorskip("torch") args = tuple((torch.tensor(arg) for arg in args)) dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = tuple((expected_jac[num] for num in argnum)) assert len(jac) == len(expected_jac) for _jac, _expected_jac in zip(jac, expected_jac): assert np.allclose(_jac, _expected_jac)
def test_jax_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and JAX.""" # Do not need the package but skip if JAX device not available pytest.importorskip("jax") dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = tuple((expected_jac[num] for num in argnum)) assert len(jac) == len(expected_jac) for _jac, _expected_jac in zip(jac, expected_jac): assert np.allclose(_jac, _expected_jac)
def test_autograd_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and Autograd.""" dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = (expected_jac[argnum[0]], ) assert len(jac) == 1 assert np.allclose(jac[0], expected_jac[0])
def test_torch_not_trainable_only(diff_method): r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Torch with ``trainable_only=False`` .""" torch = pytest.importorskip("torch") dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit_0, dev, interface="torch", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=0, trainable_only=False)( torch.tensor(a, requires_grad=True)) assert np.allclose(jac, expected_jac_not_trainable_only)
def test_autograd_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method): r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and Autograd.""" dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method) jac = classical_jacobian(qnode, argnum=argnum)(*args) expected_jac = tuple((expected_jac[num] for num in argnum)) assert len(jac) == len(expected_jac) for _jac, _expected_jac in zip(jac, expected_jac): assert np.allclose(_jac, _expected_jac)
def test_autograd_without_argnum(circuit, args, expected_jac, diff_method): r"""Test ``classical_jacobian`` with ``argnum=None`` and Autograd.""" dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method) jac = classical_jacobian(qnode)(*args) arg_shapes = [qml.math.shape(arg) for arg in args] if len(args) == 1: # For a single argument, the Jacobian is unpacked assert np.allclose(jac, expected_jac[0]) else: assert len(jac) == len(expected_jac) for _jac, _expected_jac in zip(jac, expected_jac): assert np.allclose(_jac, _expected_jac)
def test_autograd_without_argnum(circuit, args, expected_jac, diff_method): r"""Test ``classical_jacobian`` with ``argnum=None`` and Autograd.""" dev = qml.device("default.qubit", wires=2) qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method) jac = classical_jacobian(qnode)(*args) # NOTE: We use stacking to replicate qml.jacobian behaviour for scalar-only inputs if all((np.isscalar(arg) for arg in args)): expected_jac = qml.math.stack(expected_jac).T assert np.allclose(jac, expected_jac) else: # For a single argument, the Jacobian is unpacked if len(args) == 1: expected_jac = expected_jac[0] assert len(jac) == len(expected_jac) for _jac, _expected_jac in zip(jac, expected_jac): assert np.allclose(_jac, _expected_jac)