cirq.measure - accept list arguments

[pavoljuhas]

Allow passing qubits for cirq.measure and cirq.measure_each in
a single iterable argument.
Keep accepting them as variable-length arguments as before.

Fixes #2408.

[maffoo]

I think allowing arbitrarily nested lists of Qids is probably more flexibility than we want to allow. I'd suggest that we accept either Qids as varargs, or a single iterable arg containins Qids. It should be possible to use @overload types to express this to mypy in addition to the runtime handling of different args.

[pavoljuhas]

I think allowing arbitrarily nested lists of Qids is probably more flexibility than we want to allow. I'd suggest that we accept either Qids as varargs, or a single iterable arg containins Qids. ...

Sounds good to me, but Gate.on_each() already accepts a mix of varargs and possibly nested sequences.
The issue at hand (#2408) asks for a similar interface in cirq.measure().
If we leave it as is we'll get a consistent behavior in what those 2 functions can digest.

Cirq/cirq-core/cirq/ops/raw_types.py

Lines 266 to 271 in 6273826

	for target in targets:
	if isinstance(target, Qid):
	operations.append(self.on(target))
	elif isinstance(target, Iterable) and not isinstance(target, str):
	operations.extend(self.on_each(*target))
	else:

[maffoo]

The Gate.on_each implementation is a bit confused because how it behaves depends on the number of qubits the gate expects. For single-qubit gates, it allows combinations of varargs and arbitrarily-nested sequences, as you point out:

In [1]: a, b, c, d = cirq.LineQubit.range(4)

In [2]: cirq.H.on_each([[[[a]]]], b, [[c], [[d]]])
Out[2]: 
[cirq.H(cirq.LineQubit(0)),
 cirq.H(cirq.LineQubit(1)),
 cirq.H(cirq.LineQubit(2)),
 cirq.H(cirq.LineQubit(3))]

but for multi-qubit gates it only allows either varargs or a single sequence of sequences of qids:

n [3]: cirq.CZ.on_each((a, b), (c, d))
Out[3]: 
[cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(1)),
 cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(3))]

In [4]: cirq.CZ.on_each([(a, b), (c, d)])
Out[4]: 
[cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(1)),
 cirq.CZ(cirq.LineQubit(2), cirq.LineQubit(3))]

In [5]: cirq.CZ.on_each([(a, b)], [(c, d)])
...
ValueError: All values in sequence should be Qids, but got [(cirq.LineQubit(0), cirq.LineQubit(1))]

I'm not sure if this behavior for single-qubit gates was intentional, but I'd be surprised if it were. I think just accepting gate.on_each(a, b, c, d) or gate.on_each([a, b, c, d]) would be sufficient, and calling with some deeply nested list is more likely an unintentional error on the part of the caller that we should flag rather than allow.

For cirq.measure I'd suggest that rather than emulating the current behavior of Gate.on_each, we should take the more conservative approach and just allow cirq.measure(a, b, c, d) or cirq.measure([a, b, c, d]), but not arbitrarily nested sequences.

Whatever we decide to do, we should probably also modify cirq.measure_each (which produces separate single-qubit measure gates for each qubit, rather than one measure gate for all of them) in the same way, so that it stays consistent with cirq.measure.

[pavoljuhas]

@maffoo - sounds good, I will adjust it that way.
Thanks for pointing out measure_each needs the same change.

[pavoljuhas]

Ready for next review. I am not entirely sure if I got right the typing in function implementations after @overload.

[maffoo]

One minor comment, then LGTM.

[pavoljuhas]

nit: Could modify this slightly to accept exactly one arg

Done. Thank you for the mypy tip on double underscore!

[pavoljuhas]

@dabacon, @cduck, @vtomole - PTAL.

[dabacon]

LGTM