[BUG] too many subscripts in einsum with 8 qubits (default.mixed)

[ankit27kh]

Expected behavior

The circuit runs without any error.

Actual behavior

Execution fails with error:

ValueError: too many subscripts in einsum

Additional information

It is mentioned here #807 (comment) that the device can handle 23 qubits before encountering this error. This also forms the basis of this error:

ValueError: This device does not currently support computations on more than 23 wires

But, the execution fails with some circuits using just 8 qubits.
A simple example is given below.

Source code

import pennylane.numpy as np
import pennylane as qml

n = 7
dev = qml.device('default.mixed', wires=n + 1, shots=None)
state_vec = np.random.random(2 ** n)
state_vec = state_vec / np.linalg.norm(state_vec)
dm = state_vec.reshape([2 ** n, 1]) @ state_vec.reshape([2 ** n, 1]).T


@qml.qnode(dev)
def circ():
    qml.QubitDensityMatrix(dm, wires=range(n))
    qml.MultiControlledX(wires=range(n + 1))
    return qml.probs()

probs = circ()
print(probs)

Tracebacks

No response

System information

Name: PennyLane
Version: 0.28.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: 
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning, PennyLane-Lightning-GPU, PennyLane-qiskit, pennylane-qulacs
Platform info:           Linux-5.15.0-56-generic-x86_64-with-glibc2.35
Python version:          3.9.10
Numpy version:           1.22.3
Scipy version:           1.8.0
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.28.0)
- default.gaussian (PennyLane-0.28.0)
- default.mixed (PennyLane-0.28.0)
- default.qubit (PennyLane-0.28.0)
- default.qubit.autograd (PennyLane-0.28.0)
- default.qubit.jax (PennyLane-0.28.0)
- default.qubit.tf (PennyLane-0.28.0)
- default.qubit.torch (PennyLane-0.28.0)
- default.qutrit (PennyLane-0.28.0)
- null.qubit (PennyLane-0.28.0)
- qiskit.aer (PennyLane-qiskit-0.23.0)
- qiskit.basicaer (PennyLane-qiskit-0.23.0)
- qiskit.ibmq (PennyLane-qiskit-0.23.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.23.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.23.0)
- lightning.gpu (PennyLane-Lightning-GPU-0.23.0)
- qulacs.simulator (pennylane-qulacs-0.16.0)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.

[dwierichs]

This seems to be a particular issue with MultiControlledX because it truly acts on all passed wires simultaneously.
A fix for this particular operation could be to do [Hadamard(target_wire), MultiControlledZ(control_wires, target_wire), Hadamard(target_wire)] in the decomposition, which would make the only operation that really acts on all passed wires diagonal, and hence manageable. The same approach is taken e.g. by PauliRot when passed non-diagonal Pauli words.

(Note that MultiControlledZ does not exist under this name but would be needed to be decomposed further or coded up separately)

[ankit27kh]

So the actual limit for this device is seven qubits when you have gates which have not been decomposed further?
I used this with the QPE template, and the error would have come from the controlled unitaries.
Can einsum be replaced with something else? #807 (comment) here, it is mentioned that tensordot can be a replacement.
7 qubits are not nearly enough.

[dwierichs]

I'm not completely sure what the limit is, it might depend on the combination of device wire count and operation wire count.

I built a naive implementation based on tensordot:

def _apply_channel_tensordot(self, kraus, wires):
    r"""Apply a quantum channel specified by a list of Kraus operators to subsystems of the
    quantum state. For a unitary gate, there is a single Kraus operator. Implementation
    based on `tensordot` instead of `einsum`.

    Args:
        kraus (list[array]): Kraus operators
        wires (Wires): target wires
    """
    channel_wires = self.map_wires(wires)
    rho_dim = 2 * self.num_wires
    num_ch_wires = len(channel_wires)

    # Computes K^\dagger, needed for the transformation K \rho K^\dagger
    kraus_dagger = [qnp.conj(qnp.transpose(k)) for k in kraus]

    kraus = qnp.stack(kraus)
    kraus_dagger = qnp.stack(kraus_dagger)

    # Shape kraus operators
    kraus_shape = [len(kraus)] + [2] * (num_ch_wires * 2)
    kraus = qnp.cast(qnp.reshape(kraus, kraus_shape), dtype=self.C_DTYPE)
    kraus_dagger = qnp.cast(qnp.reshape(kraus_dagger, kraus_shape), dtype=self.C_DTYPE)

    # row indices of the quantum state affected by this operation
    row_wires_list = channel_wires.tolist()

    # column indices are shifted by the number of wires
    col_wires_list = [w + self.num_wires for w in row_wires_list]

    axes_left = [list(range(num_ch_wires, 2 * num_ch_wires)), row_wires_list]
    axes_right = [col_wires_list, list(range(num_ch_wires))]
    source = list(range(-num_ch_wires, 0))
    dest = col_wires_list
    self._state = qnp.sum([
            qnp.tensordot(qnp.tensordot(k, self._state, axes_left), k_d, axes_right)
            for k, k_d in zip(kraus, kraus_dagger)
        ],
        axis=0,
    )

    self._state = qnp.moveaxis(self._state, source, dest)

It passes the existing tests for apply_channel but beyond that it is untested!
My laptop allows me to allocate up to 13 qubits, and MultiControlledX works at least up until then.

[ankit27kh]

I tried this implementation, and the circuits execute with higher qubit counts. But, the results for some circuits are incorrect and do not match the original implementation.
A simple circuit that gives incorrect results:

@qml.qnode(dev)
def circ():
    qml.Hadamard(0)
    qml.SWAP([0, 2])
    qml.Hadamard(2)
    return qml.state()

After ironing out the bugs and testing for performance, this can be a solution.

[dwierichs]

Ah, that's right. I forgot to permute the axes corresponding to the row indices of the density matrix.
In any case, this should be done in a proper PR, the main effort probably is in coming up with a minimalistic but sufficient combination of test cases. The current suite neither picked up the original bug reported here, nor the one in the prototype using tensordot.

[dwierichs]

@ankit27kh If you are interested, take a look at the linked PR and the fixed prototype.
Let me know in case you find anything that's off or another use case that is not supported! Thank you! :)