Add support for > 32 qudits to cirq.sample_state_vector. Fix for #6031

[NoureldinYosri]

• When the number of qudits > 32 do the transpose operation on the flattened array.
• I expect this to be slower than np.transpose(tensor, axes) so I'm keeping the use of np.transpose when we can.
• I also rewrote the _prob in a more efficient way. The old way did a lot of extra work when reordering the axes is all that is needed.

Fix for #6031.

[rht]

This is simpler than the old code. I suppose it is also faster?

[NoureldinYosri]

yes. the main difference is what happens in the case of not measuring all qudits. In the old way it did $n$ steps each of which is slicing a matrix followed by np.sum. In the new way there is no slicing it's just a reshape since I order the indicies such that the unmeasured qudits are the last indicies so they become rows of the final matrix.

[rht]

Looking at the usage of np.transpose elsewhere in the codebase, it looks like there is another place with similar usage:

Cirq/cirq-core/cirq/sim/density_matrix_utils.py

Lines 195 to 213 in a95f009

	if len(indices) == len(qid_shape):
	# We're measuring every qudit, so no need for fancy indexing
	probs = np.abs(tensor)
	probs = np.transpose(probs, indices)
	probs = probs.reshape(-1)
	else:
	# Fancy indexing required
	meas_shape = tuple(qid_shape[i] for i in indices)
	probs = np.abs(
	[
	tensor[
	linalg.slice_for_qubits_equal_to(
	indices, big_endian_qureg_value=b, qid_shape=qid_shape
	)
	]
	for b in range(np.prod(meas_shape, dtype=np.int64))
	]
	)
	probs = np.sum(probs, axis=tuple(range(1, len(probs.shape))))

[NoureldinYosri]

this is an exact copy of the code I'm modifying 😞 I refactored this part into a seprate method which I call from both locations

[NoureldinYosri]

@rht I refactored the code so that both locations call the method.

I don't know why cirq/study/resolver_test.py::test_custom_value_not_implemented is failing but it seems to be not related to my PR since it's already failing at head

$ git branch
* master
$ pytest cirq/study/resolver_test.py -k test_custom_value_not_implemented
=========================================================================================== test session starts ===========================================================================================
platform linux -- Python 3.10.11, pytest-7.3.1, pluggy-1.0.0
Using --randomly-seed=4244270818
configfile: pyproject.toml
plugins: forked-1.6.0, cov-4.0.0, xdist-2.2.1, randomly-3.12.0, anyio-3.6.2, asyncio-0.21.0
asyncio: mode=strict
collected 47 items / 46 deselected / 1 selected                                                                                                                                                           

cirq/study/resolver_test.py x                                                                                                                                                                       [100%]

==================================================================================== 46 deselected, 1 xfailed in 0.93s ====================================================================================

[NoureldinYosri]

@rht this is now ready for merging

[rht]

Nit: I searched for the files named *_util.py; there are 2 occurrences. While *_utils.py has 7 occurrences. As such, I think it should be simulation_utils.py?

[NoureldinYosri]

this is ready for another look

[tanujkhattar]

The 32 qubit dimension constraint for numpy is more annoying that I initially thought. Another example: cirq.unitary() fails when the dimension of circuit is > 16. Note that a 2**17 x 2 ** 17 numpy array of complex64 takes ~128gb, so it's not unreasonable to expect that the protocol should work. We can add this example as part of the issue and track it in a separate PR though.

In [1]: import cirq
In [2]: q = cirq.LineQubit.range(17)
In [3]: circuit = cirq.Circuit(cirq.Z.on_each(*q))
In [4]: cirq.unitary(circuit)

[NoureldinYosri]

@tanujkhattar wouldn't computing such a unitary generally (when the unitary is not a tensor product) take a couple of hours/days? ... a single multiplication of two a $2^{17} \times 2^{17}$ matricies will take around half an hour $8^{17} \textit{operations} \approx 2251 \textit{seconds} \approx 37 \textit{minutes}$ (assumming $10^{12}$ operations per second).

---

anyway lets keep track of all the places suffering from this problem in the original issue.

---

can we merge this PR?

[NoureldinYosri]

@tanujkhattar ping