Add tfq.math.mps_1d_expectation for 1D MPS

WORKSPACE

@@ -58,11 +58,32 @@ http_archive(
     urls = ["https://github.com/bazelbuild/bazel-skylib/archive/2169ae1c374aab4a09aa90e65efe1a3aad4e279b.tar.gz"],
 )
 
+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 = "9bc32249fe83120e0d427247c4e1d4ac330bc7922fd55396b17334e1b1254f9e",
+    strip_prefix = "qsim-0.10.0",
+    urls = ["https://github.com/quantumlib/qsim/archive/refs/tags/v0.10.0.zip"],
 )
 
 # Added for crosstool in tensorflow.

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=1):
+    """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: `tf.Tensor` for an integer representing bond dimension
+            in this 1D MPS. This will create the following MPS:
+            [2, bond_dim], [bond_dim, 2, bond_dim] ... [bond_dim, 2]
+    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=tf.convert_to_tensor(bond_dim))

tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py

@@ -0,0 +1,229 @@
+# 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."""
+import numpy as np
+import tensorflow as tf
+import cirq
+import sympy
+
+from tensorflow_quantum.core.ops.math_ops import simulate_mps
+from tensorflow_quantum.python import util
+
+
+class SimulateMPS1DTest(tf.test.TestCase):
+    """Tests mps_1d."""
+
+    def test_simulate_mps_1d_inputs(self):
+        """Make sure that the mps_1d 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 1'):
+            # 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]), 0)
+
+        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]))
+
+        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)
+
+
+if __name__ == "__main__":
+    tf.test.main()