[BUG] Errors when manipulating large density matrices

[albi3ro]

Expected behavior

I expect to be able to manipulate larger density matrices, both reducing them and computing qinfo quantities from them.

Actual behavior

Einsum cannot handle that many many subscripts.

Additional information

No response

Source code

@qml.qnode(qml.device('default.qubit'))
def tester():
    qml.I(wires=range(16))
    return qml.purity(wires=(0,1))
tester()

@qml.qnode(qml.device('default.qubit'))
def t():
    qml.I(wires=range(16))
    return qml.density_matrix(wires=range(16))

t()

Tracebacks

File ~/Prog/pennylane/pennylane/measurements/purity.py:94, in PurityMP.process_state(self, state, wire_order)
     92 wire_map = dict(zip(wire_order, list(range(len(wire_order)))))
     93 indices = [wire_map[w] for w in self.wires]
---> 94 state = qml.math.dm_from_state_vector(state)
     95 return qml.math.purity(state, indices=indices, c_dtype=state.dtype)

File ~/Prog/pennylane/pennylane/math/quantum.py:519, in dm_from_state_vector(state, check_state, c_dtype)
    495 """
    496 Convenience function to compute a (full) density matrix from
    497 a state vector.
   (...)
    516 
    517 """
    518 num_wires = int(np.log2(np.shape(state)[-1]))
--> 519 return reduce_statevector(
    520     state,
    521     indices=list(range(num_wires)),
    522     check_state=check_state,
    523     c_dtype=c_dtype,
    524 )

File ~/Prog/pennylane/pennylane/math/quantum.py:474, in reduce_statevector(state, indices, check_state, c_dtype)
    468 indices2 = "".join(
    469     [ABC[num_wires + i + 1] if i in indices else ABC[i + 1] for i in consecutive_wires]
    470 )
    471 target = "".join(
    472     [ABC[i + 1] for i in sorted(indices)] + [ABC[num_wires + i + 1] for i in sorted(indices)]
    473 )
--> 474 density_matrix = einsum(
    475     f"a{indices1},a{indices2}->a{target}",
    476     state,
    477     np.conj(state),
    478     optimize="greedy",
    479 )
    481 # Return the reduced density matrix by using numpy tensor product
    482 # density_matrix = np.tensordot(state, np.conj(state), axes=(traced_system, traced_system))
    484 if batch_dim is None:

File ~/Prog/pennylane/pennylane/math/multi_dispatch.py:565, in einsum(indices, like, optimize, *operands)
    563     op2 = array(op2, like=op2[0], dtype=op2[0].dtype)
    564     return np.einsum(indices, op1, op2, like=like)
--> 565 return np.einsum(indices, *operands, like=like, optimize=optimize)

File ~/Prog/pl/lib/python3.12/site-packages/autoray/autoray.py:81, in do(fn, like, *args, **kwargs)
     79 backend = _choose_backend(fn, args, kwargs, like=like)
     80 func = get_lib_fn(backend, fn)
---> 81 return func(*args, **kwargs)

File ~/Prog/pl/lib/python3.12/site-packages/numpy/core/einsumfunc.py:1434, in einsum(out, optimize, *operands, **kwargs)
   1431         kwargs["out"] = out
   1433     # Do the contraction
-> 1434     new_view = c_einsum(einsum_str, *tmp_operands, **kwargs)
   1436 # Append new items and dereference what we can
   1437 operands.append(new_view)

ValueError: einstein sum subscripts string contains too many subscripts in the output

System information

PL master

Existing GitHub issues

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

[JerryChen97]

#Qiskit/qiskit#7700 (comment)
From another similar issue, I suppose here qml.math.quantum.reduce_statevector bruteforcely calculates the full dm by tensorproduct, which in this case somehow leads to a tensor of 33 legs violating the so called np.MAXDIM

[JerryChen97]

#Qiskit/qiskit#7700 (comment) From another similar issue, I suppose here qml.math.quantum.reduce_statevector bruteforcely calculates the full dm by tensorproduct, which in this case somehow leads to a tensor of 33 legs violating the so called np.MAXDIM

If this is true cause, I guess we won't have a quick fix until a better rdm calculator, e.g. grouping some sites together in advance.