Cirq: cirq.sample_state_vector fails when the number of qubits > 32

Description of the issue

I use the cirq.sample_state_vector to extract counts from the statevector. But it appears to not scale beyond 32 qubits due to a NumPy limitation, which hardcodes the maximum number of dimensions to be 32: https://github.com/numpy/numpy/issues/5744.

How to reproduce the issue

I’m running the following code in an NVIDIA cuQuantum Docker instance, because a CPU simulation of the code would be slow. However, this shouldn’t affect that the error still happens.

import cirq
import qsimcirq

gpu_mode = 2
options = qsimcirq.QSimOptions(gpu_mode=gpu_mode)
simulator = qsimcirq.QSimSimulator(options)

n = 33
qc = cirq.Circuit()
qs = cirq.LineQubit.range(n)
for i in range(n):
    qc.append(cirq.H(qs[i]))
    qc.append(cirq.measure(qs[i]))

result = simulator.simulate(qc)
sv = result.final_state_vector
cirq.sample_state_vector(sv, range(n))

Error message:

Traceback (most recent call last):
  File "cirq_huge_circuit.py", line 21, in <module>
    cirq.sample_state_vector(sv, range(n))
  File "/opt/conda/lib/python3.8/site-packages/cirq/sim/state_vector.py", line 214, in sample_state_vector
    probs = _probs(state_vector, indices, shape)
  File "/opt/conda/lib/python3.8/site-packages/cirq/sim/state_vector.py", line 322, in _probs
    tensor = np.reshape(state, qid_shape)
  File "<__array_function__ internals>", line 180, in reshape
  File "/opt/conda/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 298, in reshape
    return _wrapfunc(a, 'reshape', newshape, order=order)
  File "/opt/conda/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 57, in _wrapfunc
    return bound(*args, **kwds)
ValueError: maximum supported dimension for an ndarray is 32, found 33

Cirq version

cirq-core 0.14

But should apply to 1.1.x as well: https://github.com/quantumlib/Cirq/blob/8b97aa5d8f90d48e41dabaddf3c2b56ec1d2bddd/cirq-core/cirq/sim/state_vector.py#L326-L354. np.reshape is not the only operation that has the <=32 dimension constraint.