Repository: tensorflow/quantum (language: python)

Issue:
[Design] Implement MPS ops under `tfq.math`
With PRs like:
https://github.com/quantumlib/qsim/pull/374
and
https://github.com/quantumlib/qsim/pull/372
going in, It makes sense for us to start looking into upgrading to the latest qsim so that we can add the MPS ops mentioned [here](https://github.com/tensorflow/quantum/issues/380#issuecomment-724557523). In particular I think we should have something like:
1. `tfq.math.1d_expectation`
2. `tfq.math.1d_samples`
3. `tfq.math.1d_sampled_expectation`
With C++ code very similar to what we already have in https://github.com/tensorflow/quantum/blob/master/tensorflow_quantum/core/ops/tfq_simulate_expectation_op.cc (except with additional error checking to ensure the circuit is 1d and using the `MPSSimulator` from qsim etc.)


Discussion / hints:

Fix patch:
diff --git a/WORKSPACE b/WORKSPACE
index 92790b6fa..d68ca1dfe 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -3,11 +3,32 @@
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
+EIGEN_COMMIT = "12e8d57108c50d8a63605c6eb0144c838c128337"
+EIGEN_SHA256 = "f689246e342c3955af48d26ce74ac34d21b579a00675c341721a735937919b02"
+
+
+http_archive(
+    name = "eigen",
+    build_file_content = """
+cc_library(
+  name = "eigen3",
+  textual_hdrs = glob(["Eigen/**", "unsupported/**"]),
+  visibility = ["//visibility:public"],
+)
+    """,
+    sha256 = EIGEN_SHA256,
+        strip_prefix = "eigen-{commit}".format(commit = EIGEN_COMMIT),
+        urls = [
+            "https://storage.googleapis.com/mirror.tensorflow.org/gitlab.com/libeigen/eigen/-/archive/{commit}/eigen-{commit}.tar.gz".format(commit = EIGEN_COMMIT),
+            "https://gitlab.com/libeigen/eigen/-/archive/{commit}/eigen-{commit}.tar.gz".format(commit = EIGEN_COMMIT),
+        ],
+)
+
 http_archive(
     name = "qsim",
-    sha256 = "d39b9c48866ce4d6a095093ae8059444d649e851219497af99e937a74f1e9a45",
-    strip_prefix = "qsim-0.9.2-dev-20210317",
-    urls = ["https://github.com/quantumlib/qsim/archive/v0.9.2-dev+20210317.zip"],
+    sha256 = "91eb09b2697accab9c0f64d5ea6ff482f4772ce000af867670b4efaf06a35224",
+    strip_prefix = "qsim-0.10.3-dev-20211001",
+    urls = ["https://github.com/quantumlib/qsim/archive/refs/tags/v0.10.3-dev+20211001.zip"],
 )
 
 http_archive(
diff --git a/release/BUILD b/release/BUILD
index c63613e91..deeae5a42 100644
--- a/release/BUILD
+++ b/release/BUILD
@@ -39,8 +39,9 @@ sh_binary(
         "//tensorflow_quantum/core/ops:tfq_unitary_op_py",
         "//tensorflow_quantum/core/ops:tfq_utility_ops_py",
         "//tensorflow_quantum/core/ops:tfq_simulate_ops_py",
-        "//tensorflow_quantum/core/ops/math_ops:inner_product_op_py",
         "//tensorflow_quantum/core/ops/math_ops:fidelity_op_py",
+        "//tensorflow_quantum/core/ops/math_ops:inner_product_op_py",
+        "//tensorflow_quantum/core/ops/math_ops:simulate_mps_py",
         "//tensorflow_quantum/core/ops/noise:noisy_samples_op_py",
         "//tensorflow_quantum/core/ops/noise:noisy_expectation_op_py",
         "//tensorflow_quantum/core/ops/noise:noisy_sampled_expectation_op_py",
diff --git a/tensorflow_quantum/core/ops/math_ops/BUILD b/tensorflow_quantum/core/ops/math_ops/BUILD
index 6585ec5b8..f62f47311 100644
--- a/tensorflow_quantum/core/ops/math_ops/BUILD
+++ b/tensorflow_quantum/core/ops/math_ops/BUILD
@@ -17,6 +17,7 @@ cc_binary(
     srcs = [
         "tfq_inner_product.cc",
         "tfq_inner_product_grad.cc",
+        "tfq_simulate_1d_expectation.cc",
     ],
     copts = select({
         ":windows": [
@@ -65,7 +66,10 @@ cc_binary(
         "//tensorflow_quantum/core/src:adj_util",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
         "//tensorflow_quantum/core/src:util_qsim",
+        "@qsim//lib:mps_simulator",
+        "@qsim//lib:mps_statespace",
         "@qsim//lib:qsim_lib",
+        "@eigen//:eigen3",
         # tensorflow core framework
         # tensorflow core lib
         # tensorflow core protos
@@ -117,3 +121,22 @@ py_test(
         "//tensorflow_quantum/python:util",
     ],
 )
+
+py_library(
+    name = "simulate_mps_py",
+    srcs = ["simulate_mps.py"],
+    data = [":_tfq_math_ops.so"],
+    deps = [
+        "//tensorflow_quantum/core/ops:load_module",
+    ],
+)
+
+py_test(
+    name = "simulate_mps_test",
+    srcs = ["simulate_mps_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":simulate_mps_py",
+        "//tensorflow_quantum/python:util",
+    ],
+)
diff --git a/tensorflow_quantum/core/ops/math_ops/__init__.py b/tensorflow_quantum/core/ops/math_ops/__init__.py
index 14caf50eb..caf9e7648 100644
--- a/tensorflow_quantum/core/ops/math_ops/__init__.py
+++ b/tensorflow_quantum/core/ops/math_ops/__init__.py
@@ -14,5 +14,6 @@
 # ==============================================================================
 """Module for tfq.core.ops.math_ops.*"""
 
-from tensorflow_quantum.core.ops.math_ops.inner_product_op import inner_product
 from tensorflow_quantum.core.ops.math_ops.fidelity_op import fidelity
+from tensorflow_quantum.core.ops.math_ops.inner_product_op import inner_product
+from tensorflow_quantum.core.ops.math_ops.simulate_mps import mps_1d_expectation
diff --git a/tensorflow_quantum/core/ops/math_ops/simulate_mps.py b/tensorflow_quantum/core/ops/math_ops/simulate_mps.py
new file mode 100644
index 000000000..799fe38c0
--- /dev/null
+++ b/tensorflow_quantum/core/ops/math_ops/simulate_mps.py
@@ -0,0 +1,57 @@
+# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Module to register MPS simulation ops."""
+import os
+import tensorflow as tf
+from tensorflow_quantum.core.ops.load_module import load_module
+
+MATH_OP_MODULE = load_module(os.path.join("math_ops", "_tfq_math_ops.so"))
+
+
+def mps_1d_expectation(programs,
+                       symbol_names,
+                       symbol_values,
+                       pauli_sums,
+                       bond_dim=4):
+    """Calculate the expectation value of circuits wrt some operator(s)
+
+    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 specificed by programs, following the ordering
+            dictated by `symbol_names`.
+        pauli_sums: `tf.Tensor` of strings with shape [batch_size, n_ops]
+            containing the string representation of the operators that will
+            be used on all of the circuits in the expectation calculations.
+        bond_dim: Integer value used for the bond dimension during simulation.
+
+    Returns:
+        `tf.Tensor` with shape [batch_size, n_ops] that holds the
+            expectation value for each circuit with each op applied to it
+            (after resolving the corresponding parameters in).
+    """
+    return MATH_OP_MODULE.tfq_simulate_mps1d_expectation(programs,
+                                                         symbol_names,
+                                                         tf.cast(
+                                                             symbol_values,
+                                                             tf.float32),
+                                                         pauli_sums,
+                                                         bond_dim=bond_dim)
diff --git a/tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_expectation.cc b/tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_expectation.cc
new file mode 100644
index 000000000..03aecae7c
--- /dev/null
+++ b/tensorflow_quantum/core/ops/math_ops/tfq_simulate_1d_expectation.cc
@@ -0,0 +1,245 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <vector>
+
+#include "../qsim/lib/circuit.h"
+#include "../qsim/lib/formux.h"
+#include "../qsim/lib/gate_appl.h"
+#include "../qsim/lib/gates_cirq.h"
+#include "../qsim/lib/mps_simulator.h"
+#include "../qsim/lib/mps_statespace.h"
+#include "../qsim/lib/seqfor.h"
+#include "../qsim/lib/simmux.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/lib/core/error_codes.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow_quantum/core/ops/parse_context.h"
+#include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
+#include "tensorflow_quantum/core/proto/program.pb.h"
+#include "tensorflow_quantum/core/src/program_resolution.h"
+#include "tensorflow_quantum/core/src/util_qsim.h"
+
+namespace tfq {
+
+using ::tensorflow::Status;
+using ::tfq::proto::PauliSum;
+using ::tfq::proto::Program;
+
+typedef qsim::Cirq::GateCirq<float> QsimGate;
+typedef qsim::Circuit<QsimGate> QsimCircuit;
+
+class TfqSimulateMPS1DExpectationOp : public tensorflow::OpKernel {
+ public:
+  explicit TfqSimulateMPS1DExpectationOp(
+      tensorflow::OpKernelConstruction* context)
+      : OpKernel(context) {
+    // Get the bond dimension of MPS
+    // Checked that bond_dim is a positive integer >= 2 by QSim definition.
+    OP_REQUIRES_OK(context, context->GetAttr("bond_dim", &bond_dim_));
+  }
+
+  void Compute(tensorflow::OpKernelContext* context) override {
+    // TODO (mbbrough): add more dimension checks for other inputs here.
+    const int num_inputs = context->num_inputs();
+    OP_REQUIRES(context, num_inputs == 4,
+                tensorflow::errors::InvalidArgument(absl::StrCat(
+                    "Expected 4 inputs, got ", num_inputs, " inputs.")));
+
+    // Create the output Tensor.
+    const int output_dim_batch_size = context->input(0).dim_size(0);
+    const int output_dim_op_size = context->input(3).dim_size(1);
+    tensorflow::TensorShape output_shape;
+    output_shape.AddDim(output_dim_batch_size);
+    output_shape.AddDim(output_dim_op_size);
+
+    tensorflow::Tensor* output = nullptr;
+    OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
+    auto output_tensor = output->matrix<float>();
+
+    // Parse program protos.
+    std::vector<Program> programs;
+    std::vector<int> num_qubits;
+    std::vector<std::vector<PauliSum>> pauli_sums;
+
+    // TODO: remove endianness workaround introduced here:
+    // https://github.com/tensorflow/quantum/pull/610
+    // once https://github.com/quantumlib/qsim/issues/492
+    // is resolved.
+    OP_REQUIRES_OK(context,
+                   GetProgramsAndNumQubits(context, &programs, &num_qubits,
+                                           &pauli_sums, true));
+
+    std::vector<SymbolMap> maps;
+    OP_REQUIRES_OK(context, GetSymbolMaps(context, &maps));
+
+    OP_REQUIRES(context, programs.size() == maps.size(),
+                tensorflow::errors::InvalidArgument(absl::StrCat(
+                    "Number of circuits and symbol_values do not match. Got ",
+                    programs.size(), " circuits and ", maps.size(),
+                    " symbol values.")));
+
+    // Construct qsim circuits.
+    std::vector<QsimCircuit> qsim_circuits(programs.size(), QsimCircuit());
+    std::vector<QsimFusedCircuit> fused_circuits(programs.size(),
+                                                 QsimFusedCircuit({}));
+    Status parse_status = Status::OK();
+    auto p_lock = tensorflow::mutex();
+    auto construct_f = [&](int start, int end) {
+      for (int i = start; i < end; i++) {
+        Status local =
+            QsimCircuitFromProgram(programs[i], maps[i], num_qubits[i],
+                                   &qsim_circuits[i], &fused_circuits[i]);
+        // If parsing works, check MPS constraints.
+        if (local.ok()) {
+          local = CheckMPSSupported(programs[i]);
+        }
+        NESTED_FN_STATUS_SYNC(parse_status, local, p_lock);
+      }
+    };
+
+    const int num_cycles = 1000;
+    context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor(
+        output_dim_batch_size, num_cycles, construct_f);
+    OP_REQUIRES_OK(context, parse_status);
+
+    int max_num_qubits = 0;
+    for (const int num : num_qubits) {
+      max_num_qubits = std::max(max_num_qubits, num);
+    }
+
+    // Since MPS simulations have much smaller memory footprint,
+    // we do not need a ComputeLarge like we do for state vector simulation.
+    ComputeSmall(num_qubits, max_num_qubits, qsim_circuits, pauli_sums, context,
+                 &output_tensor);
+  }
+
+ private:
+  int bond_dim_;
+
+  void ComputeSmall(const std::vector<int>& num_qubits,
+                    const int max_num_qubits,
+                    const std::vector<QsimCircuit>& unfused_circuits,
+                    const std::vector<std::vector<PauliSum>>& pauli_sums,
+                    tensorflow::OpKernelContext* context,
+                    tensorflow::TTypes<float>::Matrix* output_tensor) {
+    using Simulator = qsim::mps::MPSSimulator<qsim::For, float>;
+    using StateSpace = Simulator::MPSStateSpace_;
+
+    const int output_dim_op_size = output_tensor->dimension(1);
+
+    Status compute_status = Status::OK();
+    auto c_lock = tensorflow::mutex();
+    auto DoWork = [&](int start, int end) {
+      int old_batch_index = -2;
+      int cur_batch_index = -1;
+      int largest_nq = 1;
+      int cur_op_index;
+
+      // Note: ForArgs in MPSSimulator and MPSStateState are currently unused.
+      // So, this 1 is a dummy for qsim::For.
+      Simulator sim = Simulator(1);
+      StateSpace ss = StateSpace(1);
+      auto sv = ss.Create(largest_nq, bond_dim_);
+      auto scratch = ss.Create(largest_nq, bond_dim_);
+      for (int i = start; i < end; i++) {
+        cur_batch_index = i / output_dim_op_size;
+        cur_op_index = i % output_dim_op_size;
+
+        const int nq = num_qubits[cur_batch_index];
+
+        // (#679) Just ignore empty program
+        auto unfused_gates = unfused_circuits[cur_batch_index].gates;
+        if (unfused_gates.size() == 0) {
+          (*output_tensor)(cur_batch_index, cur_op_index) = -2.0;
+          continue;
+        }
+
+        if (cur_batch_index != old_batch_index) {
+          // We've run into a new state vector we must compute.
+          // Only compute a new state vector when we have to.
+          if (nq > largest_nq) {
+            largest_nq = nq;
+            sv = ss.Create(largest_nq, bond_dim_);
+            scratch = ss.Create(largest_nq, bond_dim_);
+          }
+          // no need to update scratch_state since ComputeExpectationMPSQsim
+          // will take care of things for us.
+          ss.SetStateZero(sv);
+          for (auto gate : unfused_gates) {
+            // Can't fuse, since this might break nearest neighbor constraints.
+            qsim::ApplyGate(sim, gate, sv);
+          }
+        }
+
+        // Compute expectation values without fusing gates.
+        float exp_v = 0.0;
+        NESTED_FN_STATUS_SYNC(
+            compute_status,
+            ComputeExpectationQsim(pauli_sums[cur_batch_index][cur_op_index],
+                                   sim, ss, sv, scratch, &exp_v, false),
+            c_lock);
+        (*output_tensor)(cur_batch_index, cur_op_index) = exp_v;
+        old_batch_index = cur_batch_index;
+      }
+    };
+
+    const int64_t num_cycles =
+        200 * (int64_t(1) << static_cast<int64_t>(max_num_qubits));
+    context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor(
+        unfused_circuits.size() * output_dim_op_size, num_cycles, DoWork);
+    OP_REQUIRES_OK(context, compute_status);
+  }
+};
+
+REGISTER_KERNEL_BUILDER(
+    Name("TfqSimulateMPS1DExpectation").Device(tensorflow::DEVICE_CPU),
+    TfqSimulateMPS1DExpectationOp);
+
+REGISTER_OP("TfqSimulateMPS1DExpectation")
+    .Input("programs: string")
+    .Input("symbol_names: string")
+    .Input("symbol_values: float")
+    .Input("pauli_sums: string")
+    .Output("expectations: float")
+    .Attr("bond_dim: int >= 4 = 4")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      tensorflow::shape_inference::ShapeHandle programs_shape;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &programs_shape));
+
+      tensorflow::shape_inference::ShapeHandle symbol_names_shape;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &symbol_names_shape));
+
+      tensorflow::shape_inference::ShapeHandle symbol_values_shape;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 2, &symbol_values_shape));
+
+      tensorflow::shape_inference::ShapeHandle pauli_sums_shape;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 2, &pauli_sums_shape));
+
+      tensorflow::shape_inference::DimensionHandle output_rows =
+          c->Dim(programs_shape, 0);
+      tensorflow::shape_inference::DimensionHandle output_cols =
+          c->Dim(pauli_sums_shape, 1);
+      c->set_output(0, c->Matrix(output_rows, output_cols));
+
+      return tensorflow::Status::OK();
+    });
+
+}  // namespace tfq
diff --git a/tensorflow_quantum/core/ops/parse_context.cc b/tensorflow_quantum/core/ops/parse_context.cc
index 62404d421..dd8f941fc 100644
--- a/tensorflow_quantum/core/ops/parse_context.cc
+++ b/tensorflow_quantum/core/ops/parse_context.cc
@@ -150,7 +150,8 @@ Status GetProgramsAndProgramsToAppend(
 Status GetProgramsAndNumQubits(
     OpKernelContext* context, std::vector<Program>* programs,
     std::vector<int>* num_qubits,
-    std::vector<std::vector<PauliSum>>* p_sums /*=nullptr*/) {
+    std::vector<std::vector<PauliSum>>* p_sums /*=nullptr*/,
+    bool swap_endianness /*=false*/) {
   // 1. Parse input programs
   // 2. (Optional) Parse input PauliSums
   // 3. Convert GridQubit locations to integers.
@@ -180,10 +181,12 @@ Status GetProgramsAndNumQubits(
       Program& program = (*programs)[i];
       unsigned int this_num_qubits;
       if (p_sums) {
-        OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits,
-                                                &(p_sums->at(i))));
+        OP_REQUIRES_OK(context,
+                       ResolveQubitIds(&program, &this_num_qubits,
+                                       &(p_sums->at(i)), swap_endianness));
       } else {
-        OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits));
+        OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits,
+                                                nullptr, swap_endianness));
       }
       (*num_qubits)[i] = this_num_qubits;
     }
diff --git a/tensorflow_quantum/core/ops/parse_context.h b/tensorflow_quantum/core/ops/parse_context.h
index 4cc214385..c811b68c5 100644
--- a/tensorflow_quantum/core/ops/parse_context.h
+++ b/tensorflow_quantum/core/ops/parse_context.h
@@ -64,7 +64,8 @@ typedef absl::flat_hash_map<std::string, std::pair<int, float>> SymbolMap;
 tensorflow::Status GetProgramsAndNumQubits(
     tensorflow::OpKernelContext* context,
     std::vector<tfq::proto::Program>* programs, std::vector<int>* num_qubits,
-    std::vector<std::vector<tfq::proto::PauliSum>>* p_sums = nullptr);
+    std::vector<std::vector<tfq::proto::PauliSum>>* p_sums = nullptr,
+    bool swap_endianness = false);
 
 // Parses Cirq Program protos out of the 'circuit_specs' input Tensor. Also
 // resolves the QubitIds inside of the Program. This override also parses and
diff --git a/tensorflow_quantum/core/src/program_resolution.cc b/tensorflow_quantum/core/src/program_resolution.cc
index 4a8c39ea0..82af94d35 100644
--- a/tensorflow_quantum/core/src/program_resolution.cc
+++ b/tensorflow_quantum/core/src/program_resolution.cc
@@ -83,7 +83,8 @@ Status RegisterQubits(
 }
 
 Status ResolveQubitIds(Program* program, unsigned int* num_qubits,
-                       std::vector<PauliSum>* p_sums /*=nullptr*/) {
+                       std::vector<PauliSum>* p_sums /*=nullptr*/,
+                       bool swap_endianness /*=false*/) {
   if (program->circuit().moments().empty()) {
     // (#679) Just ignore empty program.
     // Number of qubits in empty programs is zero.
@@ -116,6 +117,10 @@ Status ResolveQubitIds(Program* program, unsigned int* num_qubits,
                                                                id_set.end());
   std::sort(ids.begin(), ids.end());
 
+  // reverse endian.
+  if (swap_endianness) {
+    std::reverse(ids.begin(), ids.end());
+  }
   absl::flat_hash_map<std::string, std::string> id_to_index;
   for (size_t i = 0; i < ids.size(); i++) {
     id_to_index[ids[i].second] = absl::StrCat(i);
@@ -128,6 +133,11 @@ Status ResolveQubitIds(Program* program, unsigned int* num_qubits,
       for (Qubit& qubit : *operation.mutable_qubits()) {
         qubit.set_id(id_to_index.at(qubit.id()));
       }
+      // reverse endian.
+      if (swap_endianness) {
+        std::reverse(operation.mutable_qubits()->begin(),
+                     operation.mutable_qubits()->end());
+      }
       // Resolve control qubit ids found in the control_qubits arg.
       absl::string_view control_qubits =
           operation.args().at("control_qubits").arg_value().string_value();
@@ -318,4 +328,59 @@ Status ResolveSymbols(
   return Status::OK();
 }
 
+Status CheckMPSSupported(const Program& program) {
+  // Check if (1) there are only 1-qubit or 2-qubit gates.
+  //          (2) each two qubit gate has neighbor qubits only.
+  //
+  // Requires: program have qubit ids resolved.
+  if (program.circuit().moments().empty()) {
+    return Status::OK();
+  }
+
+  for (auto moment : program.circuit().moments()) {
+    for (auto operation : moment.operations()) {
+      // Count the number of qubits in this operation.
+      auto qs = operation.qubits();
+      std::vector<Qubit> qubits(qs.begin(), qs.end());
+      std::vector<absl::string_view> control_ids({});
+
+      if (operation.args().find("control_qubits") != operation.args().end()) {
+        absl::string_view control_qubits =
+            operation.args().at("control_qubits").arg_value().string_value();
+        if (!control_qubits.empty()) {
+          control_ids = absl::StrSplit(control_qubits, ',');
+        }
+      }
+      const int total_num_qubits = qubits.size() + control_ids.size();
+      if (total_num_qubits > 2) {
+        return Status(
+            tensorflow::error::INVALID_ARGUMENT,
+            absl::StrCat("1D operations only support 1 and 2 qubit gates. "
+                         "Found: ",
+                         total_num_qubits, " qubit gate."));
+      }
+
+      if (total_num_qubits == 2) {
+        int j = 0;
+        std::vector<int> qids(2, -1234);
+        for (; j < qubits.size(); j++) {
+          (void)absl::SimpleAtoi(qubits[j].id(), &qids[j]);
+        }
+        for (; j < 2; j++) {
+          (void)absl::SimpleAtoi(control_ids[j], &qids[j]);
+        }
+
+        // Are the two qubits not neighbors?
+        if (std::abs((int)qids[0] - (int)qids[1]) > 1) {
+          return Status(tensorflow::error::INVALID_ARGUMENT,
+                        "A program is not in 1D topology. It contains an"
+                        " operation with qubits not neighbors each other.");
+        }
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
 }  // namespace tfq
diff --git a/tensorflow_quantum/core/src/program_resolution.h b/tensorflow_quantum/core/src/program_resolution.h
index 6889401b2..40d5760dd 100644
--- a/tensorflow_quantum/core/src/program_resolution.h
+++ b/tensorflow_quantum/core/src/program_resolution.h
@@ -37,7 +37,8 @@ namespace tfq {
 // The number of qubits in the program is recorded in `num_qubits`.
 tensorflow::Status ResolveQubitIds(
     tfq::proto::Program* program, unsigned int* num_qubits,
-    std::vector<tfq::proto::PauliSum>* p_sums = nullptr);
+    std::vector<tfq::proto::PauliSum>* p_sums = nullptr,
+    bool swap_endianness = false);
 
 // Overload which allows for strict resolution of multiple programs.
 // Will resolve GridQubits in `program` and then double check that
@@ -58,6 +59,9 @@ tensorflow::Status ResolveSymbols(
     const absl::flat_hash_map<std::string, std::pair<int, float>>& param_map,
     tfq::proto::Program* program, bool resolve_all = true);
 
+// Checks if the qubits are in 1D topology.
+tensorflow::Status CheckMPSSupported(const tfq::proto::Program& program);
+
 }  // namespace tfq
 
 #endif  // TFQ_CORE_SRC_PROGRAM_RESOLUTION
diff --git a/tensorflow_quantum/core/src/util_qsim.h b/tensorflow_quantum/core/src/util_qsim.h
index 933f1582d..f5f9c7285 100644
--- a/tensorflow_quantum/core/src/util_qsim.h
+++ b/tensorflow_quantum/core/src/util_qsim.h
@@ -142,7 +142,8 @@ tensorflow::Status ComputeExpectationQsim(const tfq::proto::PauliSum& p_sum,
                                           const SimT& sim,
                                           const StateSpaceT& ss, StateT& state,
                                           StateT& scratch,
-                                          float* expectation_value) {
+                                          float* expectation_value,
+                                          bool fuse_paulis = true) {
   // apply the gates of the pauliterms to a copy of the state vector
   // and add up expectation value term by term.
   tensorflow::Status status = tensorflow::Status::OK();
@@ -165,8 +166,14 @@ tensorflow::Status ComputeExpectationQsim(const tfq::proto::PauliSum& p_sum,
     }
     // copy from src to scratch.
     ss.Copy(state, scratch);
-    for (const qsim::GateFused<QsimGate>& fused_gate : fused_circuit) {
-      qsim::ApplyFusedGate(sim, fused_gate, scratch);
+    if (fuse_paulis) {
+      for (const qsim::GateFused<QsimGate>& fused_gate : fused_circuit) {
+        qsim::ApplyFusedGate(sim, fused_gate, scratch);
+      }
+    } else {
+      for (const auto& unfused_gate : main_circuit.gates) {
+        qsim::ApplyGate(sim, unfused_gate, scratch);
+      }
     }
 
     if (!status.ok()) {


Test patch:
diff --git a/scripts/import_test.py b/scripts/import_test.py
index 9ee4ea8d1..48a4d4b8c 100644
--- a/scripts/import_test.py
+++ b/scripts/import_test.py
@@ -37,6 +37,7 @@ def test_imports():
     # Math ops.
     _ = tfq.math.inner_product
     _ = tfq.math.fidelity
+    _ = tfq.math.mps_1d_expectation
 
     # Noisy simulation ops.
     _ = tfq.noise.expectation
diff --git a/tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py b/tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py
new file mode 100644
index 000000000..7682d1f9a
--- /dev/null
+++ b/tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py
@@ -0,0 +1,370 @@
+# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests that specifically target simulate_mps."""
+# Remove PYTHONPATH collisions for protobuf.
+# pylint: disable=wrong-import-position
+import sys
+NEW_PATH = [x for x in sys.path if 'com_google_protobuf' not in x]
+sys.path = NEW_PATH
+# pylint: enable=wrong-import-position
+
+import numpy as np
+import tensorflow as tf
+import cirq
+import cirq_google
+import sympy
+
+from tensorflow_quantum.core.ops.math_ops import simulate_mps
+from tensorflow_quantum.python import util
+
+
+class SimulateMPS1DExpectationTest(tf.test.TestCase):
+    """Tests mps_1d_expectation."""
+
+    def test_simulate_mps_1d_expectation_inputs(self):
+        """Makes sure that the op fails gracefully on bad inputs."""
+        n_qubits = 5
+        batch_size = 5
+        symbol_names = ['alpha']
+        qubits = cirq.GridQubit.rect(1, n_qubits)
+        circuit_batch = [
+            cirq.Circuit(
+                cirq.X(qubits[0])**sympy.Symbol(symbol_names[0]),
+                cirq.Z(qubits[1]),
+                cirq.CNOT(qubits[2], qubits[3]),
+                cirq.Y(qubits[4])**sympy.Symbol(symbol_names[0]),
+            ) for _ in range(batch_size)
+        ]
+        resolver_batch = [{symbol_names[0]: 0.123} for _ in range(batch_size)]
+
+        symbol_values_array = np.array(
+            [[resolver[symbol]
+              for symbol in symbol_names]
+             for resolver in resolver_batch])
+
+        pauli_sums = util.random_pauli_sums(qubits, 3, batch_size)
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'programs must be rank 1'):
+            # Circuit tensor has too many dimensions.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor([circuit_batch]), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'symbol_names must be rank 1.'):
+            # symbol_names tensor has too many dimensions.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), np.array([symbol_names]),
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'symbol_values must be rank 2.'):
+            # symbol_values_array tensor has too many dimensions.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                np.array([symbol_values_array]),
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'symbol_values must be rank 2.'):
+            # symbol_values_array tensor has too few dimensions.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array[0],
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'pauli_sums must be rank 2.'):
+            # pauli_sums tensor has too few dimensions.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array, util.convert_to_tensor(list(pauli_sums)))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'pauli_sums must be rank 2.'):
+            # pauli_sums tensor has too many dimensions.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[[x]] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'Unparseable proto'):
+            # circuit tensor has the right type but invalid values.
+            simulate_mps.mps_1d_expectation(
+                ['junk'] * batch_size, symbol_names, symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'Could not find symbol in parameter map'):
+            # symbol_names tensor has the right type but invalid values.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), ['junk'],
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'qubits not found in circuit'):
+            # pauli_sums tensor has the right type but invalid values.
+            new_qubits = [cirq.GridQubit(5, 5), cirq.GridQubit(9, 9)]
+            new_pauli_sums = util.random_pauli_sums(new_qubits, 2, batch_size)
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in new_pauli_sums]))
+
+        with self.assertRaisesRegex(TypeError, 'Cannot convert'):
+            # circuits tensor has the wrong type.
+            simulate_mps.mps_1d_expectation(
+                [1.0] * batch_size, symbol_names, symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(TypeError, 'Cannot convert'):
+            # symbol_names tensor has the wrong type.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), [0.1234],
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.UnimplementedError, ''):
+            # symbol_values tensor has the wrong type.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                [['junk']] * batch_size,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(TypeError, 'Cannot convert'):
+            # pauli_sums tensor has the wrong type.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array, [[1.0]] * batch_size)
+
+        with self.assertRaisesRegex(TypeError, 'missing'):
+            # we are missing an argument.
+            # pylint: disable=no-value-for-parameter
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array)
+            # pylint: enable=no-value-for-parameter
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'at least minimum 4'):
+            # pylint: disable=too-many-function-args
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]), 1)
+
+        with self.assertRaisesRegex(TypeError, 'Expected int'):
+            # bond_dim should be int.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]), [])
+
+        with self.assertRaisesRegex(TypeError, 'positional arguments'):
+            # pylint: disable=too-many-function-args
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]), 1, [])
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    expected_regex='do not match'):
+            # wrong op size.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums
+                                       ][:int(batch_size * 0.5)]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    expected_regex='do not match'):
+            # wrong symbol_values size.
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array[:int(batch_size * 0.5)],
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    expected_regex='cirq.Channel'):
+            # attempting to use noisy circuit.
+            noisy_circuit = cirq.Circuit(cirq.depolarize(0.3).on_each(*qubits))
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor([noisy_circuit for _ in pauli_sums]),
+                symbol_names, symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    expected_regex='not in 1D topology'):
+            # attempting to use a circuit not in 1D topology
+            # 0--1--2--3
+            #        \-4
+            circuit_not_1d = cirq.Circuit(
+                cirq.X(qubits[0])**sympy.Symbol(symbol_names[0]),
+                cirq.Z(qubits[1])**sympy.Symbol(symbol_names[0]),
+                cirq.CNOT(qubits[2], qubits[3]),
+                cirq.CNOT(qubits[2], qubits[4]),
+            )
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor([circuit_not_1d for _ in pauli_sums]),
+                symbol_names, symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    expected_regex='not in 1D topology'):
+            # attempting to use a circuit in 1D topology, which looks in 2D.
+            # 0--1
+            #  \-2-\
+            #    3--4  == 1--0--2--4--3
+            circuit_not_1d = cirq.Circuit(
+                cirq.CNOT(qubits[0], qubits[1]),
+                cirq.CNOT(qubits[0], qubits[2]),
+                cirq.CNOT(qubits[2], qubits[4]),
+                cirq.CNOT(qubits[3], qubits[4]),
+            )
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor([circuit_not_1d for _ in pauli_sums]),
+                symbol_names, symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    expected_regex='Found: 3 qubit gate'):
+            # attempting to use 3 qubit gate
+            three_qb_circuit = cirq.Circuit(
+                cirq.ISWAP(qubits[0], qubits[1]).controlled_by(qubits[2]),
+                cirq.X.on_each(*qubits))
+            simulate_mps.mps_1d_expectation(
+                util.convert_to_tensor([three_qb_circuit for _ in pauli_sums]),
+                symbol_names, symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]))
+
+        res = simulate_mps.mps_1d_expectation(
+            util.convert_to_tensor([cirq.Circuit() for _ in pauli_sums]),
+            symbol_names, symbol_values_array.astype(np.float64),
+            util.convert_to_tensor([[x] for x in pauli_sums]))
+        self.assertDTypeEqual(res, np.float32)
+
+    def test_simulate_mps_1d_expectation_simple(self):
+        """Makes sure that the op shows the same result with Cirq."""
+        n_qubits = 5
+        batch_size = 5
+        symbol_names = ['alpha']
+        qubits = cirq.GridQubit.rect(1, n_qubits)
+        circuit_batch = [
+            cirq.Circuit(
+                cirq.X(qubits[0])**sympy.Symbol(symbol_names[0]),
+                cirq.Z(qubits[1]),
+                cirq.CNOT(qubits[2], qubits[3]),
+                cirq.Y(qubits[4])**sympy.Symbol(symbol_names[0]),
+            ) for _ in range(batch_size)
+        ]
+        resolver_batch = [{symbol_names[0]: 0.123} for _ in range(batch_size)]
+
+        symbol_values_array = np.array(
+            [[resolver[symbol]
+              for symbol in symbol_names]
+             for resolver in resolver_batch])
+
+        pauli_sums = [
+            cirq.Z(qubits[0]) * cirq.X(qubits[4]) for _ in range(batch_size)
+        ]
+
+        cirq_result = [
+            cirq.Simulator().simulate_expectation_values(c, p, r)
+            for c, p, r in zip(circuit_batch, pauli_sums, resolver_batch)
+        ]
+        # Default bond_dim=4
+        mps_result = simulate_mps.mps_1d_expectation(
+            util.convert_to_tensor(circuit_batch), symbol_names,
+            symbol_values_array,
+            util.convert_to_tensor([[x] for x in pauli_sums]))
+        # Expected value of 0.349...
+        self.assertAllClose(mps_result, cirq_result)
+
+    def _make_1d_circuit(self, qubits, depth):
+        """Create a 1d ladder circuit."""
+        even_pairs = list(zip(qubits[::2], qubits[1::2]))
+        odd_pairs = list(zip(qubits[1::2], qubits[2::2]))
+        ret = cirq.Circuit()
+
+        for _ in range(depth):
+            # return ret
+            ret += [(cirq.Y(q)**np.random.random()) for q in qubits]
+            ret += [
+                cirq_google.SycamoreGate()(q0, q1)**np.random.random()
+                for q0, q1 in even_pairs
+            ]
+            ret += [(cirq.Y(q)**np.random.random()) for q in qubits]
+            ret += [
+                cirq_google.SycamoreGate()(q1, q0)**np.random.random()
+                for q0, q1 in odd_pairs
+            ]
+
+        return ret
+
+    def test_complex_equality(self):
+        """Check moderate sized 1d random circuits."""
+        batch_size = 10
+        qubits = cirq.GridQubit.rect(1, 8)
+        circuit_batch = [
+            self._make_1d_circuit(qubits, 3) for _ in range(batch_size)
+        ]
+
+        pauli_sums = [[
+            cirq.Z(qubits[0]),
+            cirq.Z(qubits[-1]),
+            cirq.Z(qubits[0]) * cirq.Z(qubits[-1]),
+            cirq.Z(qubits[0]) + cirq.Z(qubits[-1])
+        ] for _ in range(batch_size)]
+        symbol_names = []
+        resolver_batch = [{} for _ in range(batch_size)]
+
+        symbol_values_array = np.array(
+            [[resolver[symbol]
+              for symbol in symbol_names]
+             for resolver in resolver_batch])
+
+        cirq_result = [
+            cirq.Simulator().simulate_expectation_values(c, p, r)
+            for c, p, r in zip(circuit_batch, pauli_sums, resolver_batch)
+        ]
+        mps_result = simulate_mps.mps_1d_expectation(
+            util.convert_to_tensor(circuit_batch),
+            symbol_names,
+            symbol_values_array,
+            util.convert_to_tensor(pauli_sums),
+            bond_dim=32)
+        self.assertAllClose(mps_result, cirq_result, atol=1e-5)
+
+    def test_correctness_empty(self):
+        """Tests the mps op with empty circuits."""
+
+        empty_circuit = tf.raw_ops.Empty(shape=(0,), dtype=tf.string)
+        empty_symbols = tf.raw_ops.Empty(shape=(0,), dtype=tf.string)
+        empty_values = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.float32)
+        empty_paulis = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.string)
+
+        out = simulate_mps.mps_1d_expectation(empty_circuit, empty_symbols,
+                                              empty_values, empty_paulis, 32)
+
+        self.assertShapeEqual(np.zeros((0, 0)), out)
+
+
+if __name__ == "__main__":
+    tf.test.main()
diff --git a/tensorflow_quantum/core/src/program_resolution_test.cc b/tensorflow_quantum/core/src/program_resolution_test.cc
index 3cf92ae85..74a200b13 100644
--- a/tensorflow_quantum/core/src/program_resolution_test.cc
+++ b/tensorflow_quantum/core/src/program_resolution_test.cc
@@ -126,6 +126,61 @@ const std::string valid_symbol_program = R"(
   }
 )";
 
+const std::string three_qubit_op_program = R"(
+  circuit {
+    moments {
+      operations {
+        qubits {
+          id: "0_0"
+        }
+        qubits {
+          id: "0_1"
+        }
+        qubits {
+          id: "0_2"
+        }
+      }
+    }
+  }
+)";
+
+/* Qubit topology:
+  1 -- 0 -- 2
+       |
+       |
+       3
+*/
+const std::string resolved_qubit_program_not_1d = R"(
+  circuit {
+    moments {
+      operations {
+        qubits {
+          id: "0"
+        }
+        qubits {
+          id: "1"
+        }
+      }
+      operations {
+        qubits {
+          id: "0"
+        }
+        qubits {
+          id: "2"
+        }
+      }
+      operations {
+        qubits {
+          id: "0"
+        }
+        qubits {
+          id: "3"
+        }
+      }
+    }
+  }
+)";
+
 TEST(ProgramResolutionTest, ResolveQubitIdsValid) {
   Program program;
   unsigned int qubit_count;
@@ -512,5 +567,41 @@ TEST(ProgramResolutionTest, ResolveSymbolsStrictFull) {
             2.0);
 }
 
+TEST(ProgramResolutionTest, CheckMPSSupportedEmpty) {
+  Program empty;
+  EXPECT_EQ(CheckMPSSupported(empty), Status::OK());
+}
+
+TEST(ProgramResolutionTest, CheckQubitsIn1DFailedByOpWithMoreThan2Qubits) {
+  Program program_with_3qubit_op;
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      three_qubit_op_program, &program_with_3qubit_op));
+  EXPECT_EQ(CheckMPSSupported(program_with_3qubit_op),
+            Status(tensorflow::error::INVALID_ARGUMENT,
+                   "1D operations only support 1 and 2 qubit gates. "
+                   "Found: 3 qubit gate."));
+}
+
+TEST(ProgramResolutionTest,
+     CheckQubitsIn1DFailedByOpWithMoreThan2QubitsOnControlQubits) {
+  Program program_with_3qubit_op;
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      valid_program, &program_with_3qubit_op));
+  EXPECT_EQ(CheckMPSSupported(program_with_3qubit_op),
+            Status(tensorflow::error::INVALID_ARGUMENT,
+                   "1D operations only support 1 and 2 qubit gates. "
+                   "Found: 3 qubit gate."));
+}
+
+TEST(ProgramResolutionTest, CheckQubitsIn1DFailedByNot1DTopology) {
+  Program program_not_1d;
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      resolved_qubit_program_not_1d, &program_not_1d));
+  EXPECT_EQ(CheckMPSSupported(program_not_1d),
+            Status(tensorflow::error::INVALID_ARGUMENT,
+                   "A program is not in 1D topology. It contains an"
+                   " operation with qubits not neighbors each other."));
+}
+
 }  // namespace
 }  // namespace tfq