def _get_cirq_samples()

in tensorflow_quantum/core/ops/cirq_ops.py [0:0]


def _get_cirq_samples(sampler=cirq.Simulator()):
    """Get a `callable` that is a TensorFlow op that outputs circuit samples.

    Generate a TensorFlow `tf.py_function` op that when called on `tf.Tensor`s
    of circuits and parameters produces a tensor of bitstring samples from all
    the circuits.

    Args:
        sampler: Object inheriting `cirq.Sampler` to use for circuit execution.

    Returns:
        `callable` that is a Tensorflow op for taking samples.
    """
    if not isinstance(sampler, cirq.Sampler):
        raise TypeError("Passed sampler must inherit cirq.Sampler.")

    @tf.custom_gradient
    def cirq_sample(programs, symbol_names, symbol_values, num_samples):
        """Draw samples from circuits.

        Draw samples from `circuits` where each circuit will have the values in
        `symbol_values` resolved into the symbols in the circuit (with the
        ordering defined by `symbol_names`).

        ```python

        symbol_names = ['a', 'b', 'c']
        programs = tfq.convert_to_tensor(
            [cirq.Circuit(H(q0) ** sympy.Symbol('a'),
                          X(q1) ** sympy.Symbol('b'),
                          Y(q2) ** sympy.Symbol('c'))]
        )

        symbol_values = [[3,2,1]]
        n_samples = [100]

        cirq_sample(programs, symbol_names, sybmol_values, n_samples)
        ```

        Would place the values of 3 into the Symbol labeled 'a', 2 into the
        symbol labeled 'b' and 1 into the symbol labeled 'c'. Then it would
        draw 100 samples from the circuit.

        Note: In the case of circuits with varying size, all nonexistant
        samples for a particular circuit are padded with -2.

        Args:
            programs: `tf.Tensor` of strings with shape [batch_size] containing
                the string representations of the circuits to be executed.
            symbol_names: `tf.Tensor` of strings with shape [n_params], which
                is used to specify the order in which the values in
                `symbol_values` should be placed inside of the circuits in
                `programs`.
            symbol_values: `tf.Tensor` of real numbers with shape
                [batch_size, n_params] specifying parameter values to resolve
                into the circuits specified by programs, following the ordering
                dictated by `symbol_names`.
            num_samples: `tf.Tensor` with one element indicating the number of
                samples to draw.

        Returns:
            `tf.Tensor` with shape
                [batch_size, num_samples, <# qubits in largest circuit>] that
                holds samples (as boolean values) for each circuit.
        """

        def _no_grad(grad):
            raise RuntimeError(
                'Differentiation through a sampling operation is not supported.'
            )

        _input_check_helper(programs, symbol_names, symbol_values)

        if not (int(tf.size(num_samples)) == 1):
            raise ValueError("num_samples tensor must have size 1")
        if not isinstance(num_samples.dtype.as_numpy_dtype(), numbers.Integral):
            raise TypeError("num_samples tensor must be of integer type")

        serialized_programs = programs
        programs, resolvers = _batch_deserialize_helper(programs, symbol_names,
                                                        symbol_values)

        num_samples = int(num_samples.numpy())

        if isinstance(sampler, (cirq.Simulator, cirq.DensityMatrixSimulator)):
            # Only local simulators can be handled by batch_sample
            results = batch_util.batch_sample(programs, resolvers, num_samples,
                                              sampler)
            return np.array(results, dtype=np.int8), _no_grad

        # All other samplers need terminal measurement gates.
        programs = [
            p + cirq.Circuit(cirq.measure(*sorted(p.all_qubits()), key='tfq'))
            for p in programs
        ]
        max_n_qubits = max(len(p.all_qubits()) for p in programs)

        if isinstance(sampler, cirq_google.QuantumEngineSampler):
            # group samples from identical circuits to reduce communication
            # overhead. Have to keep track of the order in which things came
            # in to make sure the output is ordered correctly
            to_be_grouped = [
                (ser_prog.numpy(), resolver, index)
                for index, (
                    ser_prog,
                    resolver) in enumerate(zip(serialized_programs, resolvers))
            ]

            grouped = _group_tuples(to_be_grouped)

            # start all the necessary jobs
            results_mapping = {}
            for key, value in grouped.items():
                program = programs[value[0][1]]
                resolvers = [x[0] for x in value]
                orders = [x[1] for x in value]

                # sampler.run_sweep blocks until results are in, so go around it
                result = sampler._engine.run_sweep(
                    program=program,
                    params=resolvers,
                    repetitions=num_samples,
                    processor_ids=sampler._processor_ids,
                    gate_set=sampler._gate_set)
                results_mapping[result] = orders

            # get all results
            cirq_results = [None] * len(programs)
            for key, value in results_mapping.items():
                this_results = key.results()
                for result, index in zip(this_results, value):
                    cirq_results[index] = result

        else:
            # All other cirq.Samplers handled here.
            #TODO(zaqqwerty): replace with run_batch once Cirq #3148 is resolved
            cirq_results = []
            for p, r in zip(programs, resolvers):
                cirq_results.append(sampler.run(p, r, num_samples))

        results = []
        for r in cirq_results:
            results.append(
                tf.keras.preprocessing.sequence.pad_sequences(
                    r.measurements['tfq'],
                    maxlen=max_n_qubits,
                    dtype=np.int8,
                    value=-2,
                    padding='pre'))

        return np.array(results, dtype=np.int8), _no_grad

    @_upgrade_inputs
    def sample_generator(circuit_spec, param_names, param_values, num_samples):
        out = tf.py_function(
            func=cirq_sample,
            inp=[
                tf.stop_gradient(circuit_spec),
                tf.stop_gradient(param_names), param_values,
                tf.stop_gradient(num_samples)
            ],
            Tout=tf.int8,
        )
        out.set_shape([circuit_spec.shape[0], None, None])
        return out

    return sample_generator