Issue with QFIM for time evolved mixed state

Hi PennyLane Team,

I am attempting to calculate the Quantum Fisher Information (QFI) for a mixed state. I have prepared a Bell state and introduced a phase damping channel on both qubits. I have encoded certain theta parameters into the noisy Bell state through a time evolution governed by a Hamiltonian.
However, I am encountering an error message when attempting to calculate the QFI. In the absence of noise, the calculation is successful, but the presence of noise leads to an error. I suspect that the ApproxTimeEvolution function is causing this issue, because I realised if I just encode theta by a simple parametric rotation that solves the problem.

Could you please assist me in identifying the error and providing a solution?

import pennylane as qml
from pennylane import numpy as pnp
import numpy as np

dev = qml.device("default.mixed", wires=3)

# Parameter values
theta_values = pnp.array([np.pi / 4, np.pi / 2], requires_grad=True)  # One theta per qubit
coeff_z_values = pnp.array([[1.], [1.]],requires_grad=False)  # One Z coefficient per qubit

@qml.qnode(dev)
def bell_state_circuit(theta_values, coeff_z_values, p):
    """
    Quantum circuit to generate a 2-qubit Bell state and encode Theta by a time evolution according to a Hamiltonian.
    """
    # Apply Hadamard gate to the first qubit
    qml.Hadamard(wires=0)

    # Apply CNOT gate to entangle qubits
    qml.CNOT(wires=[0, 1])

    # Step 2: Apply Phase Damping (Dephasing) Noise
    qml.PhaseDamping(p, wires=[0])  # Apply phase damping to qubit 0
    qml.PhaseDamping(p, wires=[1])  # Apply phase damping to qubit 1

    #Encoding dynamics
    for i in range(2):
        # Starting wire of the i-th qubit
        start_wire = i

        # Get parameters for this qubit
        theta = theta_values[i]
        coeff_z = coeff_z_values[i]

        # Construct Hamiltonians 
        H_Z = qml.Hamiltonian(coeff_z, [qml.PauliZ(start_wire + j) for j in range(1)])
        
        # Apply time evolution: exp(-i theta * H_Z) and exp(-i H_X)
        qml.ApproxTimeEvolution(H_Z, theta, 1)  # Theta scales only H_Z

    # Return state

    return qml.density_matrix(wires=[0, 1])

qfi_matrix = qml.gradients.quantum_fisher(bell_state_circuit)(theta_values, coeff_z_values, 0.1)
print(qfi_matrix)

This is the full error message:

/usr/local/lib/python3.11/dist-packages/pennylane/transforms/decompose.py:337: UserWarning: Operator PhaseDamping does not define a decomposition and was not found in the target gate set. To remove this warning, add the operator name (PhaseDamping) or type (<class 'pennylane.ops.channel.PhaseDamping'>) to the gate set.
  warnings.warn(
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-103-804e5894b083> in <cell line: 0>()
     39     return qml.density_matrix(wires=[0, 1])
     40 
---> 41 qfi_matrix = qml.gradients.quantum_fisher(bell_state_circuit)(theta_values, coeff_z_values, 0.1)
     42 print(qfi_matrix)

11 frames
/usr/local/lib/python3.11/dist-packages/numpy/core/numeric.py in tensordot(a, b, axes)
   1097                 axes_b[k] += ndb
   1098     if not equal:
-> 1099         raise ValueError("shape-mismatch for sum")
   1100 
   1101     # Move the axes to sum over to the end of "a"

ValueError: shape-mismatch for sum

The PennyLane version:

Name: PennyLane
Version: 0.40.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /usr/local/lib/python3.11/dist-packages
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.1.85+-x86_64-with-glibc2.35
Python version:          3.11.11
Numpy version:           1.26.4
Scipy version:           1.14.1
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.40.0)
- default.clifford (PennyLane-0.40.0)
- default.gaussian (PennyLane-0.40.0)
- default.mixed (PennyLane-0.40.0)
- default.qubit (PennyLane-0.40.0)
- default.qutrit (PennyLane-0.40.0)
- default.qutrit.mixed (PennyLane-0.40.0)
- default.tensor (PennyLane-0.40.0)
- null.qubit (PennyLane-0.40.0)
- reference.qubit (PennyLane-0.40.0)

Hi @a.rozgonyi96,

We currently have some issues with metric_tensor (see PennyLane issue #7072), which is used for QFIM under the hood. I’m thinking this might be the cause for the behaviour you’re seeing but I’m not sure. I’ll let you know if I have any more info but for now my best guess is that you’ll need to wait for a fix.

Thank you for pointing this out. I’ll follow the discussion on it.