Merge pull request #458 from tensorflow/full_multi_qubit

Final plumbing of multi-qubit control gates.

WORKSPACE

@@ -67,9 +67,9 @@ http_archive(
 
 http_archive(
     name = "qsim",
-    sha256 = "c257c9958e7b55e85161e3d39c1d8c85148f45c82f86c69712cf815728f52faf",
-    strip_prefix = "qsim-0.7.0",
-    urls = ["https://github.com/quantumlib/qsim/archive/v0.7.0.zip"],
+    sha256 = "e2853379bde52d6277f9be4b80f54d32b3b27f7242a6c561cb34fb12d823b80e",
+    strip_prefix = "qsim-0.7.1-dev-20210126",
+    urls = ["https://github.com/quantumlib/qsim/archive/v0.7.1-dev+20210126.zip"],
 )
 
 # Added for crosstool in tensorflow.

tensorflow_quantum/core/ops/cirq_ops_test.py

@@ -296,7 +296,7 @@ def test_simulate_state_output_padding(self, op_and_sim, all_n_qubits):
                     'Simulator returned unknown type of result.' +
                     str(type(result)))
 
-        self.assertAllClose(tfq_results, manual_padded_results)
+        self.assertAllClose(tfq_results, manual_padded_results, atol=1e-5)
 
     def test_state_empty_circuit(self):
         """Test empty circuits"""

tensorflow_quantum/core/ops/math_ops/inner_product_op_test.py

@@ -250,7 +250,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size,
                 internal_wf = cirq.final_state_vector(other_batch[i][j])
                 out_arr[i][j] = np.vdot(final_wf, internal_wf)
 
-        self.assertAllClose(out, out_arr)
+        self.assertAllClose(out, out_arr, atol=1e-5)
 
     @parameterized.parameters([
         {
@@ -297,7 +297,7 @@ def test_correctness_without_symbols(self, n_qubits, batch_size,
                 internal_wf = cirq.final_state_vector(other_batch[i][j])
                 out_arr[i][j] = np.vdot(final_wf, internal_wf)
 
-        self.assertAllClose(out, out_arr)
+        self.assertAllClose(out, out_arr, atol=1e-5)
 
     def test_correctness_empty(self):
         """Test the inner product between two empty circuits."""

tensorflow_quantum/core/ops/tfq_adj_grad_op.cc

@@ -219,12 +219,30 @@ class TfqAdjointGradientOp : public tensorflow::OpKernel {
           }
 
           // Hit a parameterized gate.
-          ApplyGateDagger(
-              sim, qsim_circuits[i].gates[gradient_gates[i][j - 1].index], sv);
+          auto cur_gate =
+              qsim_circuits[i].gates[gradient_gates[i][j - 1].index];
+
+          ApplyGateDagger(sim, cur_gate, sv);
+
+          // if applicable compute control qubit mask and control value bits.
+          uint64_t mask = 0;
+          uint64_t cbits = 0;
+          for (int k = 0; k < cur_gate.controlled_by.size(); k++) {
+            uint64_t control_loc = cur_gate.controlled_by[k];
+            mask |= uint64_t{1} << control_loc;
+            cbits |= ((cur_gate.cmask >> k) & 1) << control_loc;
+          }
+
           for (int k = 0; k < gradient_gates[i][j - 1].grad_gates.size(); k++) {
             // Copy sv onto scratch2 in anticipation of non-unitary "gradient
             // gate".
             ss.Copy(sv, scratch2);
+            if (!cur_gate.controlled_by.empty()) {
+              // Gradient of controlled gattes puts zeros on diagonal which is
+              // the same as collapsing the state and then applying the
+              // non-controlled version of the gradient gate.
+              ss.BulkSetAmpl(scratch2, mask, cbits, 0, 0, true);
+            }
             qsim::ApplyGate(sim, gradient_gates[i][j - 1].grad_gates[k],
                             scratch2);
 
@@ -239,9 +257,7 @@ class TfqAdjointGradientOp : public tensorflow::OpKernel {
             (*output_tensor)(i, loc) += ss.RealInnerProduct(scratch2, scratch) +
                                         ss.RealInnerProduct(scratch, scratch2);
           }
-          ApplyGateDagger(
-              sim, qsim_circuits[i].gates[gradient_gates[i][j - 1].index],
-              scratch);
+          ApplyGateDagger(sim, cur_gate, scratch);
         }
       }
     };
@@ -315,12 +331,29 @@ class TfqAdjointGradientOp : public tensorflow::OpKernel {
         }
 
         // Hit a parameterized gate.
-        ApplyGateDagger(
-            sim, qsim_circuits[i].gates[gradient_gates[i][j - 1].index], sv);
+        // todo fix this copy.
+        auto cur_gate = qsim_circuits[i].gates[gradient_gates[i][j - 1].index];
+        ApplyGateDagger(sim, cur_gate, sv);
+
+        // if applicable compute control qubit mask and control value bits.
+        uint64_t mask = 0;
+        uint64_t cbits = 0;
+        for (int k = 0; k < cur_gate.controlled_by.size(); k++) {
+          uint64_t control_loc = cur_gate.controlled_by[k];
+          mask |= uint64_t{1} << control_loc;
+          cbits |= ((cur_gate.cmask >> k) & 1) << control_loc;
+        }
+
         for (int k = 0; k < gradient_gates[i][j - 1].grad_gates.size(); k++) {
           // Copy sv onto scratch2 in anticipation of non-unitary "gradient
           // gate".
           ss.Copy(sv, scratch2);
+          if (!cur_gate.controlled_by.empty()) {
+            // Gradient of controlled gattes puts zeros on diagonal which is
+            // the same as collapsing the state and then applying the
+            // non-controlled version of the gradient gate.
+            ss.BulkSetAmpl(scratch2, mask, cbits, 0, 0, true);
+          }
           qsim::ApplyGate(sim, gradient_gates[i][j - 1].grad_gates[k],
                           scratch2);
 
@@ -335,9 +368,7 @@ class TfqAdjointGradientOp : public tensorflow::OpKernel {
           (*output_tensor)(i, loc) += ss.RealInnerProduct(scratch2, scratch) +
                                       ss.RealInnerProduct(scratch, scratch2);
         }
-        ApplyGateDagger(sim,
-                        qsim_circuits[i].gates[gradient_gates[i][j - 1].index],
-                        scratch);
+        ApplyGateDagger(sim, cur_gate, scratch);
       }
     }
   }

tensorflow_quantum/core/ops/tfq_ps_decompose_op.cc

@@ -55,7 +55,7 @@ class TfqPsDecomposeOp : public tensorflow::OpKernel {
                                 0, context->input(0).shape(), &output));
     auto output_tensor = output->flat<tensorflow::tstring>();
 
-    const int max_buffer_moments = 3;
+    const int max_buffer_moments = 5;
 
     auto DoWork = [&](int start, int end) {
       for (int i = start; i < end; i++) {
@@ -79,21 +79,21 @@ class TfqPsDecomposeOp : public tensorflow::OpKernel {
                   phase_exponent.arg_case() == Arg::ArgCase::kSymbol) {
                 // Decompose cirq.PhasedISwapPowGate only if it is
                 // parameterized.
-                num_extra_moments = 3;
+                num_extra_moments = 5;
                 Operation new_op;
 
                 new_op = getOpForPISP(cur_op, 0, 0);
                 cur_moment.mutable_operations()->at(k) = new_op;
                 new_op = getOpForPISP(cur_op, 1, 1);
-                *cur_moment.add_operations() = new_op;
-                new_op = getOpForISP(cur_op, "XXP", exponent.symbol());
                 *temp_moment_list[0].add_operations() = new_op;
-                new_op = getOpForISP(cur_op, "YYP", exponent.symbol());
+                new_op = getOpForISP(cur_op, "XXP", exponent.symbol());
                 *temp_moment_list[1].add_operations() = new_op;
-                new_op = getOpForPISP(cur_op, 1, 0);
+                new_op = getOpForISP(cur_op, "YYP", exponent.symbol());
                 *temp_moment_list[2].add_operations() = new_op;
+                new_op = getOpForPISP(cur_op, 1, 0);
+                *temp_moment_list[3].add_operations() = new_op;
                 new_op = getOpForPISP(cur_op, 0, 1);
-                *temp_moment_list[2].add_operations() = new_op;
+                *temp_moment_list[4].add_operations() = new_op;
               }
             } else if (cur_op.gate().id() == "ISP") {
               auto exponent = cur_op_map.at("exponent");

tensorflow_quantum/core/ops/tfq_simulate_ops_test.py

@@ -319,7 +319,7 @@ def test_simulate_state_output_padding(self, all_n_qubits):
             blank_state[:wf.shape[0]] = wf
             manual_padded_results.append(blank_state)
 
-        self.assertAllClose(tfq_results, manual_padded_results)
+        self.assertAllClose(tfq_results, manual_padded_results, atol=1e-5)
 
 
 class SimulateSamplesTest(tf.test.TestCase, parameterized.TestCase):

tensorflow_quantum/core/ops/tfq_utility_ops_test.py

@@ -323,8 +323,18 @@ def test_resolve_parameters_consistency(self, n_qubits, symbol_names):
                                  expected_resolved_circuits):
             for test_m, exp_m in zip(test_c, exp_c):
                 for test_o, exp_o in zip(test_m, exp_m):
-                    tg = test_o.gate
-                    eg = exp_o.gate
+                    tg = test_o
+                    eg = exp_o
+                    if isinstance(tg, cirq.ControlledOperation):
+                        self.assertEqual(tg.controls, eg.controls)
+                        self.assertEqual(tg.control_values, eg.control_values)
+                        # Pull out controlled gate for tests later on.
+                        tg = tg.sub_operation
+                        eg = eg.sub_operation
+
+                    tg = tg.gate
+                    eg = eg.gate
+
                     self.assertEqual(type(tg), type(eg))
                     # TODO(zaqqwerty): simplify parsing when cirq build parser
                     # see core/serialize/serializer.py

tensorflow_quantum/core/serialize/serializer.py

@@ -542,14 +542,18 @@ def serialize_circuit(circuit_inp):
     # to discern controlledgates from one another otherwise. This
     # "momentary demotion" occurs with the help of the DelayedAssignmentGate.
     for i, moment in enumerate(circuit):
-        for op in moment:
-            if isinstance(op,
-                          cirq.ops.controlled_operation.ControlledOperation):
-                tfq_compatible = op.sub_operation
-                tfq_compatible._tfq_control_qubits = op.controls
-                tfq_compatible._tfq_control_values = op.control_values
-                dropped_moment = moment.without_operations_touching(op.qubits)
-                circuit[i] = dropped_moment.with_operation(tfq_compatible)
+        controlled_ops = [
+            op for op in moment if isinstance(op, cirq.ControlledOperation)
+        ]
+        new_ops = dict()
+        for op in controlled_ops:
+            tfq_compatible = op.sub_operation
+            tfq_compatible._tfq_control_qubits = op.controls
+            tfq_compatible._tfq_control_values = op.control_values
+            new_ops[op.qubits] = tfq_compatible
+
+        circuit[i] = cirq.Moment(
+            new_ops[op.qubits] if op.qubits in new_ops else op for op in moment)
 
     return SERIALIZER.serialize(circuit)
 

tensorflow_quantum/python/differentiators/gradient_test.py

@@ -158,7 +158,7 @@ def exact_grad(theta):
                 })) + [{
                     'differentiator': adjoint.Adjoint(),
                     'op': circuit_execution_ops.get_expectation_op(),
-                    'n_qubits': 5,
+                    'n_qubits': 10,
                     'n_programs': 5,
                     'n_ops': 3,
                     'symbol_names': ['a', 'b']
@@ -204,7 +204,7 @@ def test_gradients_vs_cirq_finite_difference(self, differentiator, op,
                                                     symbol_names, psums)
 
         # will this be too tight? time will tell.
-        self.assertAllClose(cirq_grads, tfq_grads, rtol=1e-2, atol=1e-2)
+        self.assertAllClose(cirq_grads, tfq_grads, rtol=2e-2, atol=2e-2)
 
     @parameterized.parameters(
         list(

tensorflow_quantum/python/util.py

@@ -26,7 +26,15 @@
 
 
 def get_supported_gates():
-    """A helper to get the gates supported by tfq."""
+    """A helper to get the gates supported by tfq.
+
+    Returns a dictionary mapping from supported gate types
+    to the number of qubits each gate operates on.
+
+    Any of these gates used in conjuction with the
+    `controlled_by` function for multi qubit control are also
+    supported.
+    """
     supported_gates = serializer.SERIALIZER.supported_gate_types()
     gate_arity_mapping_dict = dict()
     for gate in supported_gates:
@@ -46,6 +54,33 @@ def get_supported_gates():
     return gate_arity_mapping_dict
 
 
+def _apply_random_control(gate, all_qubits):
+    """Returns a random controlled version of `gate`.
+
+    Chooses a random subset s from `all_qubits` that does not intersect
+    with `gate.qubits` and returns gate.controlled_by(s). Note that
+    if no such set s can be found (size of s would be zero) then
+    `gate` is returned unchanged.
+
+    Args:
+        gate: Gate to be promoted to a controlled gate with the
+            `controlled_by` function in Cirq.
+        all_qubits: All qubits used by the circuit which `gate`
+            comes from.
+    Returns:
+        A new gate with a random subset of the set difference
+        between all_qubits and gate.qubits controlling `gate`.
+    """
+    open_qubits = set(all_qubits) - set(gate.qubits)
+    n_open = min(len(open_qubits), 3)
+    if n_open == 0:
+        # No open qubits to control. Return unmodified gate.
+        return gate
+    control_locs = random.sample(open_qubits, n_open)
+    control_values = random.choices([0, 1], k=n_open)
+    return gate.controlled_by(*control_locs, control_values=control_values)
+
+
 def random_symbol_circuit(qubits,
                           symbols,
                           n_moments=15,
@@ -64,18 +99,22 @@ def random_symbol_circuit(qubits,
     location = 0
 
     for i in range(len(circuit)):
-        if np.random.random() < p:
-            op = random.choice(list(supported_gates.keys()))
-            n_qubits = supported_gates[op]
-            locs = tuple(random.sample(qubits, n_qubits))
-            if isinstance(op, cirq.IdentityGate):
-                circuit[:i] += op.on(*locs)
-            else:
-                circuit[:i] += (op**(
-                    (np.random.random() if include_scalars else 1.0) *
-                    sympy.Symbol(random_symbols[location % len(random_symbols)])
-                )).on(*locs)
-                location += 1
+        op = random.choice(list(supported_gates.keys()))
+        n_qubits = supported_gates[op]
+        locs = tuple(random.sample(qubits, n_qubits))
+        if isinstance(op, cirq.IdentityGate):
+            circuit[:i] += op.on(*locs)
+            continue
+        working_symbol = sympy.Symbol(random_symbols[location %
+                                                     len(random_symbols)])
+        working_scalar = np.random.random() if include_scalars else 1.0
+        full_gate = (op**(working_scalar * working_symbol)).on(*locs)
+        if np.random.random() < 0.5:
+            # Add a control to this gate.
+            full_gate = _apply_random_control(full_gate, qubits)
+
+        circuit[:i] += full_gate
+        location += 1
 
     # Use the rest of the symbols
     while location < len(random_symbols):
@@ -88,12 +127,31 @@ def random_symbol_circuit(qubits,
 
 def random_circuit_resolver_batch(qubits, batch_size, n_moments=15, p=0.9):
     """Generate a batch of random circuits and symbolless resolvers."""
+    supported_gates = get_supported_gates()
     return_circuits = []
     return_resolvers = []
     for _ in range(batch_size):
-        return_circuits.append(
-            cirq.testing.random_circuit(qubits, n_moments, p,
-                                        get_supported_gates()))
+        circuit = cirq.testing.random_circuit(qubits, n_moments, p,
+                                              supported_gates)
+
+        for i in range(len(circuit)):
+            op = random.choice(list(supported_gates.keys()))
+            n_qubits = supported_gates[op]
+            if (n_qubits > len(qubits)):
+                # skip adding gates in small case.
+                continue
+            locs = tuple(random.sample(qubits, n_qubits))
+            if isinstance(op, cirq.IdentityGate):
+                circuit[:i] += op.on(*locs)
+                continue
+            full_gate = (op**np.random.random()).on(*locs)
+            if np.random.random() < 0.5:
+                # Add a control to this gate.
+                full_gate = _apply_random_control(full_gate, qubits)
+
+            circuit[:i] += full_gate
+
+        return_circuits.append(circuit)
         return_resolvers.append(cirq.ParamResolver({}))
 
     return return_circuits, return_resolvers
@@ -357,7 +415,10 @@ def get_circuit_symbols(circuit):
     for moment in circuit:
         for op in moment:
             if cirq.is_parameterized(op):
-                all_symbols |= _symbols_in_op(op.gate)
+                sub_op = op
+                if isinstance(op, cirq.ControlledOperation):
+                    sub_op = op.sub_operation
+                all_symbols |= _symbols_in_op(sub_op.gate)
     return [str(x) for x in all_symbols]
 
 

tensorflow_quantum/python/util_test.py

@@ -208,19 +208,6 @@ def test_get_circuit_symbols_all(self):
             self.assertListEqual(sorted(extracted_symbols),
                                  sorted(expected_symbols))
 
-    def test_get_circuit_symbols_error(self):
-        """Ensure that errors are reported when using unsupported ops."""
-        # TODO(mbbrough): remove this test once we reach complete parity
-        #   with cirq in terms of parametrized gate support.
-        qubits = cirq.GridQubit.rect(1, 2)
-        symbol = sympy.Symbol("u")
-        op = cirq.ZPowGate(exponent=symbol).on(qubits[0]).controlled_by(
-            qubits[1])
-        bad_circuit = cirq.Circuit(op)
-        with self.assertRaisesRegex(
-                ValueError, expected_regex="tfq.util.get_supported_gates"):
-            util.get_circuit_symbols(bad_circuit)
-
 
 class ExponentialUtilFunctionsTest(tf.test.TestCase):
     """Test that Exponential utility functions work."""