Add mps 1d Sample / SampledExpectation ops

scripts/import_test.py

@@ -38,6 +38,8 @@ def test_imports():
     _ = tfq.math.inner_product
     _ = tfq.math.fidelity
     _ = tfq.math.mps_1d_expectation
+    _ = tfq.math.mps_1d_sample
+    _ = tfq.math.mps_1d_sampled_expectation
 
     # Noisy simulation ops.
     _ = tfq.noise.expectation

tensorflow_quantum/core/ops/math_ops/BUILD

@@ -18,6 +18,8 @@ cc_binary(
         "tfq_inner_product.cc",
         "tfq_inner_product_grad.cc",
         "tfq_simulate_1d_expectation.cc",
+        "tfq_simulate_1d_samples.cc",
+        "tfq_simulate_1d_sampled_expectation.cc",
     ],
     copts = select({
         ":windows": [
@@ -127,7 +129,9 @@ py_library(
     srcs = ["simulate_mps.py"],
     data = [":_tfq_math_ops.so"],
     deps = [
+        "//tensorflow_quantum/core/ops:batch_util",
         "//tensorflow_quantum/core/ops:load_module",
+        "//tensorflow_quantum/core/ops:tfq_utility_ops_py",
     ],
 )
 

tensorflow_quantum/core/ops/math_ops/__init__.py

@@ -16,4 +16,5 @@
 
 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
+from tensorflow_quantum.core.ops.math_ops.simulate_mps import (
+    mps_1d_expectation, mps_1d_sample, mps_1d_sampled_expectation)

tensorflow_quantum/core/ops/math_ops/simulate_mps.py

@@ -16,6 +16,7 @@
 import os
 import tensorflow as tf
 from tensorflow_quantum.core.ops.load_module import load_module
+from tensorflow_quantum.core.ops import tfq_utility_ops
 
 MATH_OP_MODULE = load_module(os.path.join("math_ops", "_tfq_math_ops.so"))
 
@@ -55,3 +56,90 @@ def mps_1d_expectation(programs,
                                                              tf.float32),
                                                          pauli_sums,
                                                          bond_dim=bond_dim)
+
+
+def mps_1d_sample(programs,
+                  symbol_names,
+                  symbol_values,
+                  num_samples,
+                  bond_dim=4):
+    """Generate samples using the C++ MPS simulator.
+
+    Simulate the final state of `programs` given `symbol_values` are placed
+    inside of the symbols with the name in `symbol_names` in each circuit.
+    From there we will then sample from the final state.
+
+    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 specified by programs, following the ordering
+            dictated by `symbol_names`.
+        num_samples: `tf.Tensor` with one element indicating the number of
+            samples to draw.
+        bond_dim: Integer value used for the bond dimension during simulation.
+
+    Returns:
+        A `tf.RaggedTensor` containing the samples taken from each circuit in
+        `programs`.
+    """
+    padded_samples = MATH_OP_MODULE.tfq_simulate_mps1d_samples(
+        programs,
+        symbol_names,
+        tf.cast(symbol_values, tf.float32),
+        num_samples,
+        bond_dim=bond_dim)
+
+    return tfq_utility_ops.padded_to_ragged(padded_samples)
+
+
+def mps_1d_sampled_expectation(programs,
+                               symbol_names,
+                               symbol_values,
+                               pauli_sums,
+                               num_samples,
+                               bond_dim=4):
+    """Calculate the expectation value of circuits using samples.
+
+    Simulate the final state of `programs` given `symbol_values` are placed
+    inside of the symbols with the name in `symbol_names` in each circuit.
+    Them, sample the resulting state `num_samples` times and use these samples
+    to compute expectation values of the given `pauli_sums`.
+
+    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.
+        num_samples: `tf.Tensor` with `num_samples[i][j]` is equal to the
+            number of samples to draw in each term of `pauli_sums[i][j]`
+            when estimating the expectation. Therefore, `num_samples` must
+            have the same shape as `pauli_sums`.
+        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_sampled_expectation(
+        programs,
+        symbol_names,
+        tf.cast(symbol_values, tf.float32),
+        pauli_sums,
+        tf.cast(num_samples, dtype=tf.int32),
+        bond_dim=bond_dim)