Add tfq.math.mps_1d_expectation for 1D MPS (#610)

* initial commit

* Fix simulate_mps_test

* Add CheckQubitsIn1D

* Add Eigen & QSim MPS

* Add mps_1d op and its test

* Mike's feedback - add bond_dim as attr

* Add Attr and its tests

* Rename to mps_1d_expectation

* Fix lint and format

* Add import_test

* Fix wheel test - add mps into release/BUILD

* Mike's feedback

* WIP where Segmentation fault happens?

* Add Mike's feedback - no gate with control_qubit is allowed

* Revert control_qubit gate validation.

Both Operation with natural born controlled gate (e.g. CNOT) and
synthesized ControlledOperation are converted into the same tfq_gate_set
with control_qubit. So, we can't filter ControlledOperation in this way,
revert the previous commit.

* Fix program_resolution_test

* Fix CheckQubitsIn1D()

* Add print messages

* Add debug print

* Bump up to qsim==0.10.2

* Remove tfq::QsimFor and util_qsim.h::ComputeExpectationMPS

* Add ComputeExpectationMPSQsim in util_qsim.h

* Add more detailed debug prints

* Bump up bond_dim from 2 (default) to 4 to fix segfault error.

* Fix how to use ApplyGate in ComputeExpectationMPSQsim

* Uncomment ComputeSmall() and Remove debug outputs

* Fix format

* Mike's feedback 1 : use fuses circuit and ApplyFusedGate()

* Mike's feedback 2 : set default bond_dim to 4

* Fix format

* Mike's feedback 3 : bump up to qsim==0.10.3

* bring 1d mps expectation into working state.

* empty test.

Co-authored-by: Michael Broughton <[email protected]>

Author: jaeyoo

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(

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",

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

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",
+    ],
+)

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

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)