Shape mismatch issue when calculating QFI matrix

Hello PennyLane Community,

I’m trying to compute the Quantum Fisher Information (QFI) matrix for specific parameters of the quantum circuit using the qml.gradients.quantum_fisher function. in PennyLane.

My goal is to:
1: Prepare a Bell state.

2: Evolve it under the Hamiltonian: H = \theta_1 Z_1 + X_1 + \theta_2 Z_2 + X_2
where X_i, Z_i are the Pauli-Z and Pauli-X operators acting on qubit i.

3: Compute the QFI matrix with respect to \theta_1, \theta_2.

Here’s the code I have so far:

# System parameters
num_nodes = 2
num_qubits_per_node = 1
total_qubits = num_nodes * num_qubits_per_node
dev = qml.device('default.qubit', wires=total_qubits)

# Parameter values
theta_values = pnp.array([np.pi / 4, np.pi / 2], requires_grad=True)  # One theta per node
coeff_z_values = jnp.array([[1.], [1.]])  # One Z coefficient per node (since each node has 1 qubit)
coeff_x_values = jnp.array([[1.], [1.]])  # One X coefficient per node


@qml.qnode(dev, interface="jax")
def bell_state_circuit(theta_values, coeff_z_values, coeff_x_values):
    """
    Quantum circuit to generate a Bell state with two nodes (1 qubit per node).
    """
    # Apply Hadamard gate to the first qubit 
    qml.Hadamard(wires=0)

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

    #Encoding dynamics
    num_nodes = len(theta_values)

    for i in range(num_nodes):
        # Starting wire of the i-th node
        start_wire = i * num_qubits_per_node

        # Get parameters for this node
        theta = theta_values[i]
        coeff_z = coeff_z_values[i]
        coeff_x = coeff_x_values[i]

        # Construct separate Hamiltonians for Z and X terms
        H_Z = qml.Hamiltonian(coeff_z, [qml.PauliZ(start_wire + j) for j in range(num_qubits_per_node)])
        H_X = qml.Hamiltonian(coeff_x, [qml.PauliX(start_wire + j) for j in range(num_qubits_per_node)])

        # Apply time evolution: exp(-i theta * H_Z) and exp(-i H_X)
        qml.ApproxTimeEvolution(H_Z, theta, 1)  # Theta scales only H_Z
        qml.ApproxTimeEvolution(H_X, 1.0, 1)    # H_X evolves independently
  
    return qml.state()

qml.gradients.quantum_fisher(bell_state_circuit, argnums=[0])(theta_values, coeff_z_values, coeff_x_values)

I’m getting the following error:

/usr/local/lib/python3.11/dist-packages/pennylane/math/interface_utils.py:127: UserWarning: Contains tensors of types {'jax', 'autograd'}; dispatch will prioritize TensorFlow, PyTorch, and Jax over Autograd. Consider replacing Autograd with vanilla NumPy.
  warnings.warn(
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-63-d10aa63fb8b2> in <cell line: 0>()
     44     return qml.state()
     45 
---> 46 qml.gradients.quantum_fisher(bell_state_circuit, argnums=[0])(theta_values, coeff_z_values, coeff_x_values)

18 frames
    [... skipping hidden 18 frame]

/usr/local/lib/python3.11/dist-packages/jax/_src/lax/lax.py in _dot_general_shape_rule(lhs, rhs, dimension_numbers, precision, preferred_element_type)
   2772     msg = ("dot_general requires contracting dimensions to have the same "
   2773            "shape, got {} and {}.")
-> 2774     raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))
   2775 
   2776   return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)

TypeError: dot_general requires contracting dimensions to have the same shape, got (2,) and (4,).

Here it is the version that I use:

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.13.1
Installed devices:
- 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)
- lightning.qubit (PennyLane_Lightning-0.40.0)

Thanks for the help in advance!

Hi @a.rozgonyi96,

It look like part of the issue is that you’re mixing Jax and Numpy. I’ve modified your code to use only Numpy and it works. For this I’ve changed the coeffs from using jnp to pnp and I’ve removed the interface in the declaration of the qnode (this way the qnode will detect it automatically). Finally, I’ve removed the argnum when calculating the QFIM. Let me know if this works for you!

# System parameters
num_nodes = 2
num_qubits_per_node = 1
total_qubits = num_nodes * num_qubits_per_node
dev = qml.device('default.qubit', wires=total_qubits)

# Parameter values
theta_values = pnp.array([np.pi / 4, np.pi / 2], requires_grad=True)  # One theta per node

# Changed the coeffs to PennyLane Numpy
coeff_z_values = pnp.array([[1.], [1.]])  # One Z coefficient per node (since each node has 1 qubit)
coeff_x_values = pnp.array([[1.], [1.]])  # One X coefficient per node


@qml.qnode(dev) # removed the interface
def bell_state_circuit(theta_values, coeff_z_values, coeff_x_values):
    """
    Quantum circuit to generate a Bell state with two nodes (1 qubit per node).
    """
    # Apply Hadamard gate to the first qubit 
    qml.Hadamard(wires=0)

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

    #Encoding dynamics
    num_nodes = len(theta_values)

    for i in range(num_nodes):
        # Starting wire of the i-th node
        start_wire = i * num_qubits_per_node

        # Get parameters for this node
        theta = theta_values[i]
        coeff_z = coeff_z_values[i]
        coeff_x = coeff_x_values[i]

        # Construct separate Hamiltonians for Z and X terms
        H_Z = qml.Hamiltonian(coeff_z, [qml.PauliZ(start_wire + j) for j in range(num_qubits_per_node)])
        H_X = qml.Hamiltonian(coeff_x, [qml.PauliX(start_wire + j) for j in range(num_qubits_per_node)])

        # Apply time evolution: exp(-i theta * H_Z) and exp(-i H_X)
        qml.ApproxTimeEvolution(H_Z, theta, 1)  # Theta scales only H_Z
        qml.ApproxTimeEvolution(H_X, 1.0, 1)    # H_X evolves independently
  
    return qml.state()

# removed the argnum
qml.gradients.quantum_fisher(bell_state_circuit)(theta_values, coeff_z_values, coeff_x_values) #, argnums=[0]

Hi @CatalinaAlbornoz ,

Your help is much appreciated, based on this I could solve the problems. Thank you!

1 Like