Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] `qnode_spectrum` does not return the spectra for the correct operation
### Expected behavior

Consider the following QNode:

```python
import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(x):
    qml.RY(0.1, wires=0)
    qml.RY(0.1, wires=1)
    qml.CRX(x, wires=[0, 1])
    return qml.expval(qml.PauliZ(1))
```

We expect

```pycon
>>> qml.fourier.qnode_spectrum(circuit)(0.5)
{'x': {(): [-1.0, -0.5, 0.0, 0.5, 1.0]}}
```

since the CRX gate has 2 positive frequencies.

### Actual behavior

Instead, we get:


```pycon
>>> qml.fourier.qnode_spectrum(circuit)(0.5)
{'x': {(): [-1.0, 0.0, 1.0]}}
```

This is because the `qnode_spectrum` function is iterating through _all_ operations on the tape, rather than just the trainable operations.

As far as I can tell, this can be fixed by changing:

https://github.com/PennyLaneAI/pennylane/blob/95dfeda9a8be010b11c9e77b3a1daedcbc35b951/pennylane/fourier/qnode_spectrum.py#L423

to 
```python
op = tape.get_operation(op_idx)[0]
```
which only considers _trainable_ operations.

### Additional information

Maybe slightly unrelated, but I think

```python
tape = qml.transforms.expand_multipar(qnode.qtape)
```

should be changed to?
```
not_tape | is_measurement | has_nopar | has_gen | is_trainable
```

I noticed that when I had constant multiparameter gates in my circuit, such as `qml.Rot`, they were being expanded even when this is not required.


### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
n/a
```


### 

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

Pull request:
Fix `classical_jacobian` with JAX and operation access order in `qnode_spectrum`
**Context:**
Here we fix 2 things:
1. In JAX, the `trainable_only` argument of `get_parameters` was fixed to `False` in order to make JAX output the classical Jacobian entries belonging to gates that depend on a QNode argument but are not currently marked as trainable via `argnum` in the `classical_jacobian` call. We now assert this completeness of the Jacobian by computing it for all `argnums` instead but leaving `trainable_only=True`.
2. By using the fix proposed in #1935 we fix #1935.

**Description of the Change:**
The way the trainable operations are accessed in `qnode_spectrum` is changed so that is respects non-trainable operations in the tape correctly.

**Benefits:**
Bug fixes

**Possible Drawbacks:**
Potentially, the classical Jacobian with JAX may take small overheads with this change. However, being a bug fix, this is a negligible drawback.

**Related GitHub Issues:**
#1935 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 66dbe6f71..b1101e8a4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -396,6 +396,15 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where the classical preprocessing Jacobian
+  computed by `qml.transforms.classical_jacobian` with JAX
+  returned a reduced submatrix of the Jacobian.
+  [(#1935)](https://github.com/PennyLaneAI/pennylane/pull/1935)
+
+* Fixes a bug where the operations are not accessed in the correct order
+  in `qml.fourier.qnode_spectrum`, leading to wrong outputs.
+  [(#1935)](https://github.com/PennyLaneAI/pennylane/pull/1935)
+
 * Fixes several Pylint errors.
   [(#1951)](https://github.com/PennyLaneAI/pennylane/pull/1951)
 
diff --git a/pennylane/fourier/qnode_spectrum.py b/pennylane/fourier/qnode_spectrum.py
index 346d3d1b9..9ab404b28 100644
--- a/pennylane/fourier/qnode_spectrum.py
+++ b/pennylane/fourier/qnode_spectrum.py
@@ -399,30 +399,30 @@ def qnode_spectrum(qnode, encoding_args=None, argnum=None, decimals=8, validatio
                     f"The return_type {m.return_type.value} is not supported as it likely does "
                     "not admit a Fourier spectrum."
                 )
-        class_jacs = jac_fn(*args, **kwargs)
+        cjacs = jac_fn(*args, **kwargs)
         spectra = {}
         tape = qml.transforms.expand_multipar(qnode.qtape)
         par_info = tape._par_info
 
         # Iterate over jacobians per argument
-        for jac_idx, class_jac in enumerate(class_jacs):
+        for jac_idx, cjac in enumerate(cjacs):
             # Obtain argument name for the jacobian index
             arg_name = arg_name_map[jac_idx]
             # Extract requested parameter indices for the current argument
             if encoding_args[arg_name] is Ellipsis:
                 # If no index for this argument is specified, request all parameters within
                 # the argument (Recall () is a valid index for scalar-valued arguments here)
-                requested_par_ids = set(product(*(range(sh) for sh in class_jac.shape[1:])))
+                requested_par_ids = set(product(*(range(sh) for sh in cjac.shape[1:])))
             else:
                 requested_par_ids = set(encoding_args[arg_name])
             # Each requested parameter at least "contributes" as a constant
             _spectra = {par_idx: {0} for par_idx in requested_par_ids}
 
             # Iterate over the axis of the current Jacobian that corresponds to the tape operations
-            for op_idx, jac_of_op in enumerate(np.round(class_jac, decimals=decimals)):
+            for op_idx, jac_of_op in enumerate(np.round(cjac, decimals=decimals)):
                 op = par_info[op_idx]["op"]
                 # Find parameters that both were requested and feed into the operation
-                if len(class_jac.shape) == 1:
+                if len(cjac.shape) == 1:
                     # Scalar argument, only axis of Jacobian is for operations
                     if np.isclose(jac_of_op, 0.0, atol=atol, rtol=0):
                         continue
diff --git a/pennylane/transforms/classical_jacobian.py b/pennylane/transforms/classical_jacobian.py
index a96a1dcce..4cf62c500 100644
--- a/pennylane/transforms/classical_jacobian.py
+++ b/pennylane/transforms/classical_jacobian.py
@@ -19,7 +19,7 @@ import pennylane as qml
 from pennylane import numpy as np
 
 
-def classical_jacobian(qnode, argnum=None, expand_fn=None):
+def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
     r"""Returns a function to extract the Jacobian
     matrix of the classical part of a QNode.
 
@@ -137,7 +137,6 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None):
 
     def classical_preprocessing(*args, **kwargs):
         """Returns the trainable gate parameters for a given QNode input."""
-        trainable_only = kwargs.pop("_trainable_only", True)
         kwargs.pop("shots", None)
         qnode.construct(args, kwargs)
         tape = qnode.qtape
@@ -185,7 +184,14 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None):
         argnum = 0 if argnum is None else argnum
 
         def _jacobian(*args, **kwargs):
-            kwargs["_trainable_only"] = False
+            if trainable_only:
+                _argnum = list(range(len(args)))
+                full_jac = jax.jacobian(classical_preprocessing, argnums=_argnum)(*args, **kwargs)
+                if np.isscalar(argnum):
+                    return full_jac[argnum]
+
+                return tuple(full_jac[i] for i in argnum)
+
             return jax.jacobian(classical_preprocessing, argnums=argnum)(*args, **kwargs)
 
         return _jacobian


Test patch:
diff --git a/tests/fourier/test_qnode_spectrum.py b/tests/fourier/test_qnode_spectrum.py
index 40ebd3fcc..edb0d1ddb 100644
--- a/tests/fourier/test_qnode_spectrum.py
+++ b/tests/fourier/test_qnode_spectrum.py
@@ -32,6 +32,7 @@ def circuit_0(a):
 def circuit_1(a, b):
     qml.RZ(-a / 3, wires=0)
     qml.RX(a / 5, wires=1)
+    qml.CRZ(0.3, wires=[1, 0])
     qml.CNOT(wires=[0, 1])
     qml.RY(b * 2, wires=1)
     qml.RZ(-b, wires=1)
@@ -395,8 +396,8 @@ class TestAutograd:
         """Test that the spectra of a circuit is calculated correctly
         in the autograd interface."""
 
-        x = pnp.array([1.0, 2.0, 3.0])
-        w = pnp.array([[-1, -2, -3], [-4, -5, -6]], dtype=float)
+        x = pnp.array([1.0, 2.0, 3.0], requires_grad=True)
+        w = pnp.array([[-1, -2, -3], [-4, -5, -6]], dtype=float, requires_grad=True)
 
         dev = qml.device("default.qubit", wires=3)
         qnode = qml.QNode(circuit, dev, interface="autograd")
@@ -421,7 +422,7 @@ class TestTorch:
         in the torch interface."""
 
         torch = pytest.importorskip("torch")
-        x = torch.tensor([1.0, 2.0, 3.0])
+        x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
         w = torch.tensor([[-1, -2, -3], [-4, -5, -6]], dtype=float)
 
         dev = qml.device("default.qubit", wires=3)
diff --git a/tests/transforms/test_classical_jacobian.py b/tests/transforms/test_classical_jacobian.py
index 4548ef4f4..0f9906f8d 100644
--- a/tests/transforms/test_classical_jacobian.py
+++ b/tests/transforms/test_classical_jacobian.py
@@ -26,6 +26,7 @@ z = np.array([2.1, -0.3, 0.62, 0.89])
 
 
 def circuit_0(a):
+    qml.RZ(0.2, wires=0)
     [qml.RX(a, wires=0) for i in range(4)]
     return qml.expval(qml.PauliZ(0))
 
@@ -70,7 +71,7 @@ def circuit_5(x, y, z):
 
 
 circuits = [circuit_0, circuit_1, circuit_2, circuit_3, circuit_4, circuit_5]
-args = [(a,), (a, b), (x,), (x, y), (x, y), (x, y, z)]
+all_args = [(a,), (a, b), (x,), (x, y), (x, y), (x, y, z)]
 interfaces = ["jax", "autograd", "tf", "torch"]
 
 class_jacs = [
@@ -112,7 +113,7 @@ class_jacs = [
 
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
-@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, all_args, class_jacs))
 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)
@@ -133,7 +134,7 @@ def test_autograd_without_argnum(circuit, args, expected_jac, diff_method):
 
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
-@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, all_args, class_jacs))
 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
@@ -147,7 +148,7 @@ def test_jax_without_argnum(circuit, args, expected_jac, diff_method):
 
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
-@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, all_args, class_jacs))
 def test_tf_without_argnum(circuit, args, expected_jac, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=None`` and Tensorflow."""
     tf = pytest.importorskip("tensorflow")
@@ -162,7 +163,7 @@ def test_tf_without_argnum(circuit, args, expected_jac, diff_method):
 
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
-@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, all_args, class_jacs))
 def test_torch_without_argnum(circuit, args, expected_jac, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=None`` and Torch."""
     torch = pytest.importorskip("torch")
@@ -181,7 +182,7 @@ scalar_argnum = [0, 1, 0, 1, 0, 1]
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, scalar_argnum)
 )
 def test_autograd_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Autograd."""
@@ -194,7 +195,7 @@ def test_autograd_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_m
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, scalar_argnum)
 )
 def test_jax_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and JAX."""
@@ -209,7 +210,7 @@ def test_jax_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, scalar_argnum)
 )
 def test_tf_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and TensorFlow."""
@@ -224,7 +225,7 @@ def test_tf_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method)
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, scalar_argnum)
 )
 def test_torch_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Torch."""
@@ -242,7 +243,7 @@ single_list_argnum = [[0], [1], [0], [1], [0], [2]]
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, single_list_argnum)
 )
 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."""
@@ -256,7 +257,7 @@ def test_autograd_with_single_list_argnum(circuit, args, expected_jac, argnum, d
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, single_list_argnum)
 )
 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."""
@@ -272,7 +273,7 @@ def test_jax_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_m
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, single_list_argnum)
 )
 def test_tf_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and TensorFlow."""
@@ -288,7 +289,7 @@ def test_tf_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_me
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, single_list_argnum)
 )
 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."""
@@ -307,7 +308,7 @@ sequence_argnum = [[0], [0, 1], (0,), [0, 1], (0, 1), {0, 2}]
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, sequence_argnum)
 )
 def test_autograd_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and Autograd."""
@@ -322,7 +323,7 @@ def test_autograd_with_sequence_argnum(circuit, args, expected_jac, argnum, diff
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, sequence_argnum)
 )
 def test_jax_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and JAX."""
@@ -339,7 +340,7 @@ def test_jax_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_meth
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, sequence_argnum)
 )
 def test_tf_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and TensorFlow."""
@@ -356,7 +357,7 @@ def test_tf_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_metho
 
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
-    "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
+    "circuit, args, expected_jac, argnum", zip(circuits, all_args, class_jacs, sequence_argnum)
 )
 def test_torch_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and Torch."""
@@ -369,3 +370,52 @@ def test_torch_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_me
     assert len(jac) == len(expected_jac)
     for _jac, _expected_jac in zip(jac, expected_jac):
         assert np.allclose(_jac, _expected_jac)
+
+
+expected_jac_not_trainable_only = np.array([0.0, 1.0, 1.0, 1.0, 1.0])
+
+
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+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)
+
+
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+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)
+
+
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+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)
+
+
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+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)