Unexpected behavior when using batch_jacobian with multiple inputs/outputs in quantum-classical neural network

username: ziv_chen
cooked: <p>I’m implementing a neural network that includes quantum layers (using PennyLane’s pytorch’s interface) to solve ODEs. I want to encode several points at once and get several results of the ODE (one result per corresponding input) at the same run.<br>
Currently I’m running a toy model where I try to get <code>u(x)=sin(x)</code> according to the loss function: <code>du_dx - cos(x)</code>.</p>
<p>The network structure is:</p>
<details>
<summary>
network structure and quantum circuit</summary>
<pre><code class="lang-auto">class QuantumNetworkBaseAngleClass(nn.Module):
    def __init__(self, n_qubits=8, n_layers=2, x_scale_factor=1.0, x_scale_bias=0.0):
        super().__init__()
        self.n_qubits = n_qubits 
        self.n_layers = n_layers
        self.weight_shapes = None
        self.x_scale_factor = x_scale_factor
        self.x_scale_bias = x_scale_bias
        self.q_layer, self.qnode = self.create_hybrid_network()

        # Initialize the weights of the quantum layer
        for param in self.q_layer.parameters():
            if param.requires_grad:  # Only initialize trainable parameters
                nn.init.uniform_(param, 0.0, torch.pi)

    def create_hybrid_network(self):
        dev = qml.device("default.qubit", wires=self.n_qubits)
        layer = qml.StronglyEntanglingLayers     
        shape = layer.shape(n_layers=self.n_layers, n_wires=self.n_qubits)
        self.weight_shapes = {"weights": shape}

        @qml.qnode(dev, interface="torch", diff_method="backprop", differentiable=True, max_diff=2)
        def qnode(inputs, weights):
            scaled_inputs = (inputs - self.x_scale_bias) * self.x_scale_factor
            qml.AngleEmbedding(scaled_inputs, wires=range(self.n_qubits))
            layer(weights, wires=range(self.n_qubits))
            return [qml.expval(qml.PauliZ(wires=i)) for i in range(self.n_qubits)]
        return qml.qnn.TorchLayer(qnode, self.weight_shapes), qnode

    def forward(self, inputs):
        output = self.q_layer(inputs)
        return output

class QuantumNetworkAngleClass(nn.Module):
    def __init__(self, n_qubits=8, n_layers=2, x_scale_factor=1.0, x_scale_bias=0.0):
        super().__init__()
        self.fc_in = nn.Linear(n_qubits, n_qubits)
        self.activation = nn.Tanh()
        self.q_layer = QuantumNetworkBaseAngleClass(n_qubits=n_qubits, n_layers=n_layers,
                                                    x_scale_factor=x_scale_factor, x_scale_bias=x_scale_bias)
        self.fc_out = nn.Linear(n_qubits, n_qubits)

    def forward(self, inputs):
        out_fc_in = self.fc_in(inputs.unsqueeze(0))
        out_activation = self.activation(out_fc_in)
        out_q_layer = self.q_layer(out_activation)
        out_fc_out = self.fc_out(out_q_layer).squeeze(0).squeeze(0)
        return out_fc_out
</code></pre>
</details>
<p>The gradients:</p>
<details>
<summary>
gradients</summary>
<pre><code class="lang-auto">        for epoch in range(self.epochs):
            u_out_arr = []
            d_u_dx_arr = []
            loss_single_batch_arr = []
            loss_single_batch_no_beta_arr = []
            for collocation_points in self.dataloader:
                self.opt.zero_grad(set_to_none=False)  # check if adding False makes any difference
                collocation_points = collocation_points.squeeze(-1)
                u_out = self.model(collocation_points)
                u_out_arr.append(u_out)  # for plotting and derivative calculation check
                d_u_dx = torch.zeros(self.n_points)
                for i in range(self.num_batches):
                    jac_func = jacobian(self.model, collocation_points[i], create_graph=True, strict=True)
                    d_u_dx[i*self.batch_size:(i+1)*self.batch_size] = torch.diag(jac_func, 0)
                d_u_dx_arr.append(d_u_dx)
                loss, loss_no_beta = self.loss_fn(collocation_points, u=u_out, d_u_dx=d_u_dx,
                                                 target_fn=self.target_fn_dict)
                loss.backward()
                loss_single_batch_arr.append(loss)
                loss_single_batch_no_beta_arr.append(loss_no_beta)

                self.opt.step()
</code></pre>
</details>
<p>When computing derivatives using jacobian (only the diagonal because I need each output’s derivative with respect to the corresponding input):</p>
<ul>
<li>With <code>in_out_size=1</code>: derivatives correctly correspond to spatial derivatives</li>
<li>With <code>in_out_size=n_qubits</code>: derivatives don’t match expected spatial derivatives</li>
</ul>
<p>Question: Why does increasing input/output dimensions affect the derivative computation?</p>
<p>The same behavior is observed when I change the quantum layer to <code>nn.Linear</code> or when I use <code>torch.gradient</code>, so I don’t think it’s due to the quantum circuit or the kind of derivative.</p>
<p>results for the two cases:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff" title="results_pytorch"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff_2_690x388.jpeg" alt="results_pytorch" data-base62-sha1="fH69LmCggKEVzjRUinT5AiZhFEz" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff_2_690x388.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff_2_1035x582.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/d/6dffd4799040f9c82e2216c6f8cb0d6ff91d3cff.jpeg 2x" data-dominant-color="FBFBFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">results_pytorch</span><span class="informations">1280×720 65.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Thanks in advance!</p>

username: Timmi315
cooked: <p>Hi ziv_chen,</p>
<p>I’m working on a similar project and am having quite similar problems. Have you verified, that the qnode receives the right format of the data?</p>
<p>I’ve been trying to use the native implementation of a qml.nn.TorchLayer for a while now but always get unexpected behaviour:</p>
<ul>
<li>I suppose you are using batched inputs, so your QuantumNetworkAngleClass receives a tensor containing multiple input vectors.</li>
<li>As I understand the pennylane docs and am reading your code, the qnode is written to only receive one input vector. Calling the qnode for eatch vector in the batch should be handled by the torch framewok (e.g. in this case your nn.Module implementation of QuantumNetworkBaseAngleClass with the torch quantum layer self.q_layer).</li>
<li>Now, when I had similar problems in my project, I started to read out the inputs my qnode gets and voila, torch passes the whole batch of inputs to my qnode (instead of one input at a time), resulting in the qnode returning very random values.</li>
<li>Since the returns from the qnode are batched again by the torch framwork (putting all return values in a tensor and reformatting it to have the “correct” shape), the rest of the model continues using the values as intended and does not detect any errors.</li>
</ul>
<p>As of now, I have not found a way for torch to only call the qnode once for every input, but at least that might be where your problems come from (since if you are only using one input and one qubit, the batched input tensor is processed one at a time and thus no mismatch of the return values can mess up the result).</p>
<p><strong> → Maybe try printing the output of the qnode and the output of self.q_layer and see if theres a mismatch, as was the case with my code</strong></p>
<p>(with batch_size of 2 and 8 qubits I want my qnode to receive and return a tensor with shape (8,) but torch passes a tensor of shape (2,8) and thus my qnode returns shape (2,) 8 times which is coverted into a tensor of shape (2,8) after the qnode execution.)</p>
<p>Good luck and I’d love to hear back!</p>

username: ziv_chen
cooked: <p>Hi <a class="mention" href="/u/timmi315">@Timmi315</a>,</p>
<p>Thanks for the thorough reply!<br>
I have tried to do what you wrote and I have few findings and a few follow-up questions.<br>
(TL;DR: still couldn’t get it to work…)</p>
<p>Findings:</p>
<ol>
<li>I used the same amount of points and qubits so I don’t have batches (4 for the example). The inputs had the shape and type of <code>Tensor: (4,)</code>. It is very hard to understand the output of the qnode itself. That’s because when I print from within the circuit, it just says <code>[expval(Z(0)), expval(Z(1)), expval(Z(2)), expval(Z(3))]</code>.<br>
I dug deeper into the code using the debugger and the return type and shape of the <code>execute</code> within the <code>qnode</code> is <code>tuple:(4)</code> (each element in the tuple is a tensor), but later it gets converted to a list of tensors, one tensor per measurement still within the <code>qnode</code>.<br>
When torch evaluates the <code>qnode</code> (using <code>self._evaluate_qnode(inputs)</code>) it receives a <code>Tensor: (4,)</code>.</li>
<li>when I use the circuit as it is with the same seed, inputs and weights I get the same output for the <code>qnode</code>, where the type is a list of tensors.</li>
<li>even without batches, the results still have the same behavior as before (even if I use more points, like12).</li>
</ol>
<p>Follow-up questions:</p>
<ol>
<li>If I would compare to your example, do you mean that my <code>qnode</code> returns shape of (1,) 4 times and I get (4,) after conversion or that it just does not apply in my case?</li>
<li>do you have any Idea of how to progress from this according to what I’ve found?</li>
</ol>
<p>Thanks!</p>
<p>By the way, the possibility to send batches to the <code>qnode</code> exists, it is called parameter broadcasting (<a href="https://pennylane.ai/blog/2022/10/how-to-execute-quantum-circuits-in-collections-and-batches" class="inline-onebox" rel="noopener nofollow ugc">How to execute quantum circuits in collections and batches | PennyLane Blog</a>) and it can speed things up quite dramatically.</p>