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)
