Finish support for LineQubit everywhere. (#641)

MichaelBroughton · web-flow · commit 81df6b31bde6 · 2021-11-17T15:09:04.000-08:00
* format.

* Jae feedback.
diff --git a/tensorflow_quantum/core/ops/circuit_execution_ops_test.py b/tensorflow_quantum/core/ops/circuit_execution_ops_test.py
@@ -175,7 +175,8 @@ def test_supported_gates_consistent(self, op_and_sim):
         """Ensure that supported gates are consistent across backends."""
         op = op_and_sim[0]
         sim = op_and_sim[1]
-        qubits = cirq.GridQubit.rect(1, 5)
+        # mix qubit types.
+        qubits = cirq.GridQubit.rect(1, 4) + [cirq.LineQubit(10)]
         circuit_batch = []
 
         gate_ref = util.get_supported_gates()
@@ -345,7 +346,7 @@ def test_analytical_expectation(self, op_and_sim, n_qubits, symbol_names,
         op = op_and_sim[0]
         sim = op_and_sim[1]
 
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits - 1) + [cirq.GridQubit(0, 0)]
         circuit_batch, resolver_batch = \
             util.random_symbol_circuit_resolver_batch(
                 qubits, symbol_names, BATCH_SIZE)
diff --git a/tensorflow_quantum/core/ops/math_ops/fidelity_op_test.py b/tensorflow_quantum/core/ops/math_ops/fidelity_op_test.py
@@ -113,7 +113,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size,
     def test_correctness_without_symbols(self, n_qubits, batch_size,
                                          inner_dim_size):
         """Tests that inner_product works without symbols."""
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
         circuit_batch, _ = \
             util.random_circuit_resolver_batch(
                 qubits, batch_size)
diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_grad_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_grad_test.py
@@ -235,7 +235,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size,
         """Tests that inner_product works with symbols."""
         symbol_names = ['alpha', 'beta', 'gamma']
         n_params = len(symbol_names)
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
         circuit_batch, resolver_batch = \
           util.random_symbol_circuit_resolver_batch(
               qubits, symbol_names, batch_size)
diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_op_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_op_test.py
@@ -232,7 +232,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size,
                                       inner_dim_size):
         """Tests that inner_product works with symbols."""
         symbol_names = ['alpha', 'beta', 'gamma']
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
         circuit_batch, resolver_batch = \
             util.random_symbol_circuit_resolver_batch(
                 qubits, symbol_names, batch_size)
diff --git a/tensorflow_quantum/core/ops/noise/noisy_expectation_op_test.py b/tensorflow_quantum/core/ops/noise/noisy_expectation_op_test.py
@@ -273,7 +273,7 @@ def test_single_channel(self, channel):
         symbol_names = []
         batch_size = 5
         n_qubits = 6
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
 
         circuit_batch, resolver_batch = \
             util.random_circuit_resolver_batch(
diff --git a/tensorflow_quantum/core/ops/noise/noisy_sampled_expectation_op_test.py b/tensorflow_quantum/core/ops/noise/noisy_sampled_expectation_op_test.py
@@ -273,7 +273,7 @@ def test_single_channel(self, channel):
         symbol_names = []
         batch_size = 5
         n_qubits = 6
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
 
         circuit_batch, resolver_batch = \
             util.random_circuit_resolver_batch(
diff --git a/tensorflow_quantum/core/ops/noise/noisy_samples_op_test.py b/tensorflow_quantum/core/ops/noise/noisy_samples_op_test.py
@@ -164,7 +164,7 @@ def test_simulate_samples_inputs(self):
     def test_simulate_consistency(self, batch_size, n_qubits, noisy):
         """Test consistency with batch_util.py simulation."""
         symbol_names = ['alpha', 'beta']
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
 
         circuit_batch, resolver_batch = \
             util.random_symbol_circuit_resolver_batch(
diff --git a/tensorflow_quantum/core/ops/tfq_adj_grad_op_test.py b/tensorflow_quantum/core/ops/tfq_adj_grad_op_test.py
@@ -359,7 +359,7 @@ def test_calculate_adj_grad_simple_case_single(self):
         n_qubits = 2
         batch_size = 1
         symbol_names = ['alpha', 'beta', 'gamma']
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
         circuit_batch, resolver_batch = \
         [cirq.Circuit(cirq.X(qubits[0]) ** sympy.Symbol('alpha'),
             cirq.Y(qubits[1]) ** sympy.Symbol('alpha'),
diff --git a/tensorflow_quantum/core/ops/tfq_ps_util_ops_test.py b/tensorflow_quantum/core/ops/tfq_ps_util_ops_test.py
@@ -259,7 +259,7 @@ def test_moment_preservation(self):
         """Test Moment-structure preservation."""
         t = sympy.Symbol('t')
         r = sympy.Symbol('r')
-        qubits = cirq.GridQubit.rect(1, 6)
+        qubits = cirq.LineQubit.range(6)
         circuit_batch = [
             cirq.Circuit(
                 cirq.Moment([cirq.H(q) for q in qubits]),
@@ -471,7 +471,7 @@ def test_weight_coefficient(self):
 
     def test_simple_pad(self):
         """Test simple padding."""
-        bit = cirq.GridQubit(0, 0)
+        bit = cirq.LineQubit(1)
         circuit = cirq.Circuit(
             cirq.X(bit)**sympy.Symbol('alpha'),
             cirq.Y(bit)**sympy.Symbol('alpha'),
@@ -720,7 +720,7 @@ def test_error(self):
 
     def test_many_values(self):
         """Ensure that padding with few symbols and many values works."""
-        bit = cirq.GridQubit(0, 0)
+        bit = cirq.LineQubit(1)
         circuits = [
             cirq.Circuit(
                 cirq.X(bit)**(sympy.Symbol('alpha') * 2.0),
diff --git a/tensorflow_quantum/core/ops/tfq_unitary_op_test.py b/tensorflow_quantum/core/ops/tfq_unitary_op_test.py
@@ -170,7 +170,7 @@ def test_calculate_unitary_consistency_symbol_free(self, n_qubits,
                                                        unitary_op):
         """Test calculate_unitary works without symbols."""
         unitary_op = tfq_unitary_op.get_unitary_op()
-        qubits = cirq.GridQubit.rect(1, n_qubits)
+        qubits = cirq.LineQubit.range(n_qubits)
         circuit_batch, _ = util.random_circuit_resolver_batch(qubits, 25)
 
         tfq_results = unitary_op(util.convert_to_tensor(circuit_batch), [],
diff --git a/tensorflow_quantum/core/ops/tfq_utility_ops_test.py b/tensorflow_quantum/core/ops/tfq_utility_ops_test.py
@@ -90,7 +90,7 @@ def test_append_circuit(self, max_n_bits, symbols, n_circuits):
         base_circuits = []
         circuits_to_append = []
         qubits = cirq.GridQubit.rect(1, max_n_bits)
-        other_qubits = cirq.GridQubit.rect(2, max_n_bits)
+        other_qubits = cirq.GridQubit.rect(2, max_n_bits) + [cirq.LineQubit(10)]
 
         base_circuits, _ = util.random_symbol_circuit_resolver_batch(
             qubits,
diff --git a/tensorflow_quantum/core/src/program_resolution.cc b/tensorflow_quantum/core/src/program_resolution.cc
@@ -40,10 +40,16 @@ using tfq::proto::PauliTerm;
 using tfq::proto::Program;
 using tfq::proto::Qubit;
 
+inline absl::string_view IntMaxStr() {
+  static constexpr char kMaxVal[] = "2147483647";
+  return kMaxVal;
+}
+
 Status RegisterQubits(
     const std::string& qb_string,
     absl::flat_hash_set<std::pair<std::pair<int, int>, std::string>>* id_set) {
   // Inserts qubits found in qb_string into id_set.
+  // Supported GridQubit wire formats and line qubit wire formats.
 
   if (qb_string.empty()) {
     return Status::OK();  // no control-default value specified in serializer.py
@@ -52,7 +58,11 @@ Status RegisterQubits(
   const std::vector<absl::string_view> qb_list = absl::StrSplit(qb_string, ',');
   for (auto qb : qb_list) {
     int r, c;
-    const std::vector<absl::string_view> splits = absl::StrSplit(qb, '_');
+    std::vector<absl::string_view> splits = absl::StrSplit(qb, '_');
+    if (splits.size() == 1) {  // Pad the front of linequbit with INTMAX.
+      splits.insert(splits.begin(), IntMaxStr());
+    }
+
     if (splits.size() != 2) {
       return Status(tensorflow::error::INVALID_ARGUMENT,
                     absl::StrCat("Unable to parse qubit: ", qb));
diff --git a/tensorflow_quantum/core/src/program_resolution_test.cc b/tensorflow_quantum/core/src/program_resolution_test.cc
@@ -54,6 +54,29 @@ const std::string valid_program = R"(
   }
 )";
 
+const std::string valid_line_program = R"(
+  circuit {
+    moments {
+      operations {
+        args {
+          key: "control_qubits"
+          value {
+            arg_value {
+              string_value: "0_1"
+            }
+          }
+        }
+        qubits {
+          id: "1"
+        }
+        qubits {
+          id: "2"
+        }
+      }
+    }
+  }
+)";
+
 const std::string valid_psum = R"(
   terms {
     coefficient_real: 1.0
@@ -123,6 +146,26 @@ TEST(ProgramResolutionTest, ResolveQubitIdsValid) {
             "0");
 }
 
+TEST(ProgramResolutionTest, ResolveQubitIdsValidLine) {
+  Program program;
+  unsigned int qubit_count;
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(valid_line_program,
+                                                            &program));
+
+  EXPECT_EQ(ResolveQubitIds(&program, &qubit_count), Status::OK());
+  EXPECT_EQ(qubit_count, 3);
+  EXPECT_EQ(program.circuit().moments(0).operations(0).qubits(0).id(), "1");
+  EXPECT_EQ(program.circuit().moments(0).operations(0).qubits(1).id(), "2");
+  EXPECT_EQ(program.circuit()
+                .moments(0)
+                .operations(0)
+                .args()
+                .at("control_qubits")
+                .arg_value()
+                .string_value(),
+            "0");
+}
+
 TEST(ProgramResolutionTest, ResolveQubitIdsInvalidControlQubit) {
   Program program;
   unsigned int qubit_count;
diff --git a/tensorflow_quantum/python/util_test.py b/tensorflow_quantum/python/util_test.py
@@ -39,7 +39,7 @@ def _exponential(theta, op):
     return np.eye(op_mat.shape[0]) * np.cos(theta) - 1j * op_mat * np.sin(theta)
 
 
-BITS = list(cirq.GridQubit.rect(1, 10))
+BITS = list(cirq.GridQubit.rect(1, 10) + cirq.LineQubit.range(2))
 
 
 def _items_to_tensorize():
