Repository: PennyLaneAI/pennylane

Issue #1991: [BUG] Metric tensor with only some of multiple QNode argument does not work in Torch
Created: 2021-12-06T18:21:26Z
Expected behavior
In the following QNode, we have two arguments:
dev = qml.device("default.qubit", wires=4)

@qml.qnode(dev, interface="torch")
def circuit(x, y):
    qml.Rot(*x[0], wires=0)
    qml.Rot(*x[1], wires=1)
    qml.Rot(*x[2], wires=2)
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    qml.CNOT(wires=[2, 0])
    qml.RY(y[0], wires=0)
    qml.RY(y[1], wires=1)
    qml.RY(y[0], wires=2)
    return qml.expval(qml.PauliZ(0)@qml.PauliZ(2))
    
x = torch.tensor([[0.2, 0.4, -0.1], [-2.1, 0.5, -0.2], [0.1, 0.7, -0.6]], dtype=torch.float64)
y = torch.tensor([1.3, 0.2], requires_grad=True, dtype=torch.float64)
The expected behaviour for the metric_tensor now would be a single array output with the metric tensor w.r.t y, namely
>>> qml.metric_tensor(circuit)(x, y)
tensor([[ 0.2550, -0.0709],
            [-0.0709,  0.2495]], dtype=torch.float64)
Actual behavior
The computation raises an error (see below).
This is as far as I can tell because the classical Jacobian is computed for all QNode in arguments in Torch. That is, without argnum as keyword, the QNode wrapper does not know which of the classical Jacobians to use to contract with the tape metric tensor(s). Another reason to implement #1880 soonish.
Indeed, when setting hybrid=False in metric_tensor, we get the correct metric tensor for the three gates that are controlled by entries of the QNode argument y:
>>> qml.metric_tensor(circuit, hybrid=False)(x, y)
tensor([[ 0.2498,  0.0032, -0.1224],
            [ 0.0032,  0.2495, -0.0741],
            [-0.1224, -0.0741,  0.2500]], dtype=torch.float64)
Additional information
No response
Source code
No response
Tracebacks
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_37953/3753152100.py in <module>
     19 # y = pnp.array([1.3, 0.2], requires_grad=True)
     20 # print(qml.adjoint_metric_tensor(circuit)(y))
---> 21 print(qml.metric_tensor(circuit)(x, y))
     22 # circuit(x, y)
     23 # circuit.qtape.get_parameters()

~/venvs/xanadu/lib/python3.8/site-packages/PennyLane-0.20.0.dev0-py3.8.egg/pennylane/transforms/metric_tensor.py in wrapper(*args, **kwargs)
    333                 for c in cjac:
    334                     if c is not None:
--> 335                         _mt = qml.math.tensordot(mt, c, axes=[[-1], [0]])
    336                         _mt = qml.math.tensordot(c, _mt, axes=[[0], [0]])
    337                         metric_tensors.append(_mt)

~/venvs/xanadu/lib/python3.8/site-packages/PennyLane-0.20.0.dev0-py3.8.egg/pennylane/math/multi_dispatch.py in tensordot(tensor1, tensor2, axes)
    280     """
    281     interface = _multi_dispatch([tensor1, tensor2])
--> 282     return np.tensordot(tensor1, tensor2, axes=axes, like=interface)
    283 
    284 

~/venvs/xanadu/lib/python3.8/site-packages/autoray/autoray.py in do(fn, like, *args, **kwargs)
     82         backend = infer_backend(like)
     83 
---> 84     return get_lib_fn(backend, fn)(*args, **kwargs)
     85 
     86 

~/venvs/xanadu/lib/python3.8/site-packages/PennyLane-0.20.0.dev0-py3.8.egg/pennylane/math/single_dispatch.py in _tensordot_torch(tensor1, tensor2, axes)
    461     if not semantic_version.match(">=1.10.0", torch.__version__) and axes == 0:
    462         return torch.outer(tensor1, tensor2)
--> 463     return torch.tensordot(tensor1, tensor2, axes)
    464 
    465 

~/venvs/xanadu/lib/python3.8/site-packages/torch/functional.py in tensordot(a, b, dims, out)
    930 
    931     if out is None:
--> 932         return _VF.tensordot(a, b, dims_a, dims_b)  # type: ignore[attr-defined]
    933     else:
    934         return _VF.tensordot(a, b, dims_a, dims_b, out=out)  # type: ignore[attr-defined]

RuntimeError: contracted dimensions need to match, but first has size 3 in dim -1 and second has size 12 in dim 0
System information
Name: PennyLane
Version: 0.20.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: /home/david/venvs/xanadu/lib/python3.8/site-packages/PennyLane-0.20.0.dev0-py3.8.egg
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, scipy, semantic-version, toml
Required-by: pennylane-qulacs, PennyLane-Qchem, PennyLane-Lightning
Platform info:           Linux-5.4.0-91-generic-x86_64-with-glibc2.29
Python version:          3.8.10
Numpy version:           1.21.0
Scipy version:           1.7.1
Installed devices:
- qulacs.simulator (pennylane-qulacs-0.16.0)
- lightning.qubit (PennyLane-Lightning-0.18.0)
- default.gaussian (PennyLane-0.20.0.dev0)
- default.mixed (PennyLane-0.20.0.dev0)
- default.qubit (PennyLane-0.20.0.dev0)
- default.qubit.autograd (PennyLane-0.20.0.dev0)
- default.qubit.jax (PennyLane-0.20.0.dev0)
- default.qubit.tf (PennyLane-0.20.0.dev0)
- default.qubit.torch (PennyLane-0.20.0.dev0)
- default.tensor (PennyLane-0.20.0.dev0)
- default.tensor.tf (PennyLane-0.20.0.dev0)


 I have searched existing GitHub issues to make sure the issue does not already exist.

Duplicate issue #2018: [BUG] Classical jacobian is computed for non-trainable parameters too with Torch
Created: 2021-12-13T14:11:30Z
Expected behavior
Using qml.transforms.classical_jacobian by passing argnum=None with Torch interface computes the classical jacobian wrt. the trainable parameters.
import pennylane as qml
from pennylane import numpy as np
import torch

dev = qml.device('default.qubit', wires=3)

@qml.qnode(dev, interface='torch')
def circuit(x, y, z, a):
    qml.RX(qml.math.sin(x), wires=0)
    qml.CNOT(wires=[0, 1])
    qml.RY(y ** 2, wires=1)
    qml.RZ(1 / z, wires=1)
    qml.RX(3*a, wires= 0)
    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

jac_fn = qml.transforms.classical_jacobian(circuit, argnum=None)

x = torch.tensor(0.1, requires_grad=True)
y = torch.tensor(-2.5, requires_grad=True)
z = torch.tensor(0.71, requires_grad=True)

a = torch.tensor(0.1, requires_grad=False)
jac_fn(x, y, z, a)
Actual behavior
We get results wrt. all parameters (non-trainable included):
(tensor([0.9950, 0.0000, 0.0000, 0.0000]),
 tensor([-0., -5., -0., -0.]),
 tensor([-0.0000, -0.0000, -1.9837, -0.0000]),
 tensor([0., 0., 0., 3.]))

Additional information
The same snippet using autograd computes the classical jacobian only wrt. trainable params:
import pennylane as qml
from pennylane import numpy as np
import torch

dev = qml.device('default.qubit', wires=3)

@qml.qnode(dev, interface='autograd')
def circuit(x, y, z, a):
    qml.RX(qml.math.sin(x), wires=0)
    qml.CNOT(wires=[0, 1])
    qml.RY(y ** 2, wires=1)
    qml.RZ(1 / z, wires=1)
    qml.RX(3*a, wires= 0)
    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

jac_fn = qml.transforms.classical_jacobian(circuit, argnum=None)

x = np.array(0.1, requires_grad=True)
y = np.array(-2.5, requires_grad=True)
z = np.array(0.71, requires_grad=True)

a = np.array(0.1, requires_grad=False)
jac_fn(x, y, z, a)
tensor([[ 0.99500417, -0.        , -0.        ],
        [ 0.        , -5.        , -0.        ],
        [ 0.        , -0.        , -1.98373339]], requires_grad=True)

Source code
No response
Tracebacks
No response
System information
Name: PennyLane
Version: 0.21.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/antal/xanadu/pennylane
Requires: numpy, scipy, networkx, autograd, toml, appdirs, semantic_version, autoray, cachetools, pennylane-lightning
Required-by: PennyLane-Cirq, PennyLane-Orquestra, pennylane-qulacs, amazon-braket-pennylane-plugin, PennyLane-Honeywell, PennyLane-qiskit, PennyLane-AQT, PennyLane-PQ, PennyLane-Forest, PennyLane-qsharp, PennyLane-Lightning, PennyLane-Qchem, PennyLane-IonQ, PennyLane-SF
Platform info:           Linux-5.11.0-41-generic-x86_64-with-glibc2.10
Python version:          3.8.5
Numpy version:           1.21.4
Scipy version:           1.7.3
Installed devices:
- cirq.mixedsimulator (PennyLane-Cirq-0.19.0)
- cirq.pasqal (PennyLane-Cirq-0.19.0)
- cirq.qsim (PennyLane-Cirq-0.19.0)
- cirq.qsimh (PennyLane-Cirq-0.19.0)
- cirq.simulator (PennyLane-Cirq-0.19.0)
- orquestra.forest (PennyLane-Orquestra-0.15.0)
- orquestra.ibmq (PennyLane-Orquestra-0.15.0)
- orquestra.qiskit (PennyLane-Orquestra-0.15.0)
- orquestra.qulacs (PennyLane-Orquestra-0.15.0)
- qulacs.simulator (pennylane-qulacs-0.17.0.dev0)
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- honeywell.hqs (PennyLane-Honeywell-0.16.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.18.0.dev0)
- aqt.noisy_sim (PennyLane-AQT-0.18.0)
- aqt.sim (PennyLane-AQT-0.18.0)
- projectq.classical (PennyLane-PQ-0.18.0.dev0)
- projectq.ibm (PennyLane-PQ-0.18.0.dev0)
- projectq.simulator (PennyLane-PQ-0.18.0.dev0)
- forest.numpy_wavefunction (PennyLane-Forest-0.18.0.dev0)
- forest.qvm (PennyLane-Forest-0.18.0.dev0)
- forest.wavefunction (PennyLane-Forest-0.18.0.dev0)
- microsoft.QuantumSimulator (PennyLane-qsharp-0.19.0)
- lightning.qubit (PennyLane-Lightning-0.20.0.dev0)
- ionq.qpu (PennyLane-IonQ-0.17.0.dev0)
- ionq.simulator (PennyLane-IonQ-0.17.0.dev0)
- default.gaussian (PennyLane-0.21.0.dev0)
- default.mixed (PennyLane-0.21.0.dev0)
- default.qubit (PennyLane-0.21.0.dev0)
- default.qubit.autograd (PennyLane-0.21.0.dev0)
- default.qubit.jax (PennyLane-0.21.0.dev0)
- default.qubit.tf (PennyLane-0.21.0.dev0)
- default.qubit.torch (PennyLane-0.21.0.dev0)
- strawberryfields.fock (PennyLane-SF-0.20.0.dev0)
- strawberryfields.gaussian (PennyLane-SF-0.20.0.dev0)
- strawberryfields.gbs (PennyLane-SF-0.20.0.dev0)
- strawberryfields.remote (PennyLane-SF-0.20.0.dev0)
- strawberryfields.tf (PennyLane-SF-0.20.0.dev0)


### 

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Closing comment (2021-12-13T14:14:00Z): Note that this is related to #1991