Shape mismatch issue when calculating QFI matrix

PennyLane Help
1

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!

2

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]
3

Hi @CatalinaAlbornoz ,

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

1 Like
4

Dear @CatalinaAlbornoz ,

Sorry for revisiting this issue, but I’d appreciate your insight once again. For now, I’m using many Trotter steps in the ApproxTimeEvolution function, which is significantly slowing down the computation of gradient.quantum_fisher. To improve performance, I think I need to integrate JAX. However, since I’m combining PennyLane (using qfim) with JAX (for the circuit), I’m still encountering compatibility issues. Do you have any suggestions on how to use them together effectively?

Thank you for your help again!

5

Hi!
We have demos about Using JAX with PennyLane and How to optimize a QML model using JAX and Optax. I believe you will find these useful and they could be a good source to give your some pointers for your application. Take a look at the PennyLane documentation as well.

I hope these are useful resources. Let us know of any further questions you might have.

1 Like
6

Dear @daniela.murcillo and @CatalinaAlbornoz ,

Thank you! I tried using the documentation and the demos, but unfortunately I couldn’t fix the error in the way that would be optimal for me.

The built-in PennyLane QFI transform (qml.gradients.quantum_fisher) isn’t working smoothly with my JAX-enabled QNode because of the internal handling of gradients.

I think because the quantum_fisher was originally designed for PennyLane’s native QNode and array types, mixing in the JAX interface can lead to kinds of mismatches in the handled parameter shapes.

When using multiple input parameters (e.g. separate arrays for different types of coefficients) with a JAX-enabled QNode, the gradients computed for each input can end up having different shapes. For example, if I understand correctly, one set of parameters might produce a gradient with shape (2,), while another yields (4,), causing a mismatch when PennyLane’s quantum_fisher transform attempts to combine them via tensor contractions.

This is my code:

# System parameters
num_qubit = 2
dev = qml.device('default.qubit', wires=2)

# Parameter values
coeff_z_values = jnp.array([jnp.pi, jnp.pi])  # One Z coefficient per qubit
coeff_x_values = jnp.array([1., 1.])  # One X coefficient per qubit

@qml.qnode(dev, interface='jax')
def bell_state_circuit(coeff_z_values, coeff_x_values):
    """
    Quantum circuit to generate a Bell state and add extra parametric gates.
    """
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0, 1])

    #Encoding dynamics
    for i in range(num_qubit):
        start_wire = i 
        # Get parameters for this qubit
        coeff_z = coeff_z_values[i]
        coeff_x = coeff_x_values[i]

        # Adding Z and X terms
        qml.RZ(coeff_z, start_wire)
        qml.RX(coeff_x, start_wire)
         
    return qml.state()

qml.gradients.quantum_fisher(bell_state_circuit)(coeff_z_values, coeff_x_values)

And the error message I received:

/usr/local/lib/python3.11/dist-packages/pennylane/workflow/resolution.py:78: RuntimeWarning: PennyLane is currently not compatible with versions of JAX > 0.4.28. You have version 0.4.33 installed.
  warn(
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-38-0e20f42f1fb4> in <cell line: 0>()
----> 1 qml.gradients.quantum_fisher(bell_state_circuit)(coeff_z_values, coeff_x_values)

17 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,).

However for me it seems, when bundling all parameters into a single flat input vector—and then splitting them appropriately inside the circuit—ensures that the gradient information is consistently structured. This consistency avoids the shape mismatch problem when computing the quantum Fisher information. Essentially, using one unified parameter vector leads to uniformly shaped derivatives, which are compatible with the quantum_fisher transform.

Modified working code:

# System parameters
num_qubit = 2
dev = qml.device('default.qubit', wires=num_qubit)

@qml.qnode(dev, interface='jax')
def bell_state_circuit(params):
    """
    Quantum circuit that generates a Bell state and applies extra parametric gates.
    
    The input `params` is a 1D array with:
      - The first num_qubit elements corresponding to the Z rotation angles.
      - The next num_qubit elements corresponding to the X rotation angles.
    """
    # Split the parameters: first half are Z coefficients, second half are X coefficients.
    coeff_z_values = params[:num_qubit]
    coeff_x_values = params[num_qubit:]
    
    # Generate the Bell state.
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0, 1])
    
    # Apply parametric rotations on each qubit.
    for i in range(num_qubit):
        qml.RZ(coeff_z_values[i], wires=i)
        qml.RX(coeff_x_values[i], wires=i)
         
    return qml.state()

qml.gradients.quantum_fisher(bell_state_circuit)(params)

This, however, makes things a bit more difficult for me, because I would prefer to handle the variables separately. Moreover, it would be sufficient for me to compute the QFIM with respect to one of the parameters, computing it with both is unnecessary.

Could you, please, help me figure out how to avoid the mismatch error in the first code version? Thank you for your suggestions in advance!

The qml version I used:

Name: PennyLane
Version: 0.41.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.12
Numpy version:           2.0.2
Scipy version:           1.14.1
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.41.0)
- default.clifford (PennyLane-0.41.0)
- default.gaussian (PennyLane-0.41.0)
- default.mixed (PennyLane-0.41.0)
- default.qubit (PennyLane-0.41.0)
- default.qutrit (PennyLane-0.41.0)
- default.qutrit.mixed (PennyLane-0.41.0)
- default.tensor (PennyLane-0.41.0)
- null.qubit (PennyLane-0.41.0)
- reference.qubit (PennyLane-0.41.0)