Inner product (tensorflow#387)

* Add tfq.math.inner_product.

* format.

* updated release/BUILD with op.

* added missing test.

* CI numpy 1.18.5 version pin.

* add test_output=all

* CI expt 2: revert load_module.py

* CI expt 3: name conlifcts on math? -> math_ops.

* CI expt 3 build fix.

* remove duplicate program size check.

release/BUILD

@@ -12,6 +12,7 @@ sh_binary(
         # Core module.
         "//tensorflow_quantum/core:__init__.py",
         "//tensorflow_quantum/core/ops:__init__.py",
+        "//tensorflow_quantum/core/ops/math_ops:__init__.py",
         "//tensorflow_quantum/core/proto:__init__.py",
         "//tensorflow_quantum/core/serialize:__init__.py",
 
@@ -37,6 +38,7 @@ 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/serialize:serializer",
         "//tensorflow_quantum/datasets:cluster_state",
         "//tensorflow_quantum/datasets:spin_system",

scripts/import_test.py

@@ -34,6 +34,9 @@ def test_imports():
     _ = tfq.padded_to_ragged2d
     _ = tfq.resolve_parameters
 
+    # Math ops.
+    _ = tfq.math.inner_product
+
     # Util functions.
     _ = tfq.convert_to_tensor
     _ = tfq.get_quantum_concurrent_op_mode

tensorflow_quantum/__init__.py

@@ -21,6 +21,9 @@
                                      get_unitary_op, padded_to_ragged,
                                      padded_to_ragged2d, resolve_parameters)
 
+# Import math ops.
+from tensorflow_quantum.core import math_ops as math
+
 # Re-label python module as layers module.
 import tensorflow_quantum.python.layers as layers
 

tensorflow_quantum/core/__init__.py

@@ -24,3 +24,6 @@
 # Special case for append op which we didn't name well.
 from tensorflow_quantum.core.ops import \
     tfq_append_circuit as append_circuit
+
+# Import math ops.
+from tensorflow_quantum.core.ops import math_ops

tensorflow_quantum/core/ops/__init__.py

@@ -24,3 +24,6 @@
                                                          padded_to_ragged2d,
                                                          resolve_parameters,
                                                          tfq_append_circuit)
+
+# Import math_ops.
+from tensorflow_quantum.core.ops import math_ops

tensorflow_quantum/core/ops/math_ops/BUILD

@@ -0,0 +1,84 @@
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])
+
+# Export for the PIP package.
+exports_files(["__init__.py"])
+
+config_setting(
+    name = "windows",
+    constraint_values = ["@bazel_tools//platforms:windows"],
+)
+
+cc_binary(
+    name = "_tfq_math_ops.so",
+    srcs = [
+        "tfq_inner_product.cc",
+    ],
+    copts = select({
+        ":windows": [
+            "/D__CLANG_SUPPORT_DYN_ANNOTATION__",
+            "/D_USE_MATH_DEFINES",
+            "/DEIGEN_MPL2_ONLY",
+            "/DEIGEN_MAX_ALIGN_BYTES=64",
+            "/DEIGEN_HAS_TYPE_TRAITS=0",
+            "/DTF_USE_SNAPPY",
+            "/showIncludes",
+            "/MD",
+            "/O2",
+            "/DNDEBUG",
+            "/w",
+            "-DWIN32_LEAN_AND_MEAN",
+            "-DNOGDI",
+            "/d2ReducedOptimizeHugeFunctions",
+            "/arch:AVX",
+            "/std:c++14",
+            "-DTENSORFLOW_MONOLITHIC_BUILD",
+            "/DPLATFORM_WINDOWS",
+            "/DEIGEN_HAS_C99_MATH",
+            "/DTENSORFLOW_USE_EIGEN_THREADPOOL",
+            "/DEIGEN_AVOID_STL_ARRAY",
+            "/Iexternal/gemmlowp",
+            "/wd4018",
+            "/wd4577",
+            "/DNOGDI",
+            "/UTF_COMPILE_LIBRARY",
+        ],
+        "//conditions:default": [
+            "-pthread",
+            "-std=c++11",
+            "-D_GLIBCXX_USE_CXX11_ABI=0",
+        ],
+    }),
+    features = select({
+        ":windows": ["windows_export_all_symbols"],
+        "//conditions:default": [],
+    }),
+    linkshared = 1,
+    deps = [
+        "//tensorflow_quantum/core/ops:parse_context",
+        "//tensorflow_quantum/core/ops:tfq_simulate_utils",
+        "//tensorflow_quantum/core/src:util_qsim",
+        "//tensorflow_quantum/core/src:circuit_parser_qsim",
+        "@qsim//lib:qsim_lib",
+    ],
+)
+
+py_library(
+    name = "inner_product_op_py",
+    srcs = ["inner_product_op.py"],
+    data = [":_tfq_math_ops.so"],
+    deps = [
+        "//tensorflow_quantum/core/ops:load_module",
+    ],
+)
+
+py_test(
+    name = "inner_product_op_test",
+    srcs = ["inner_product_op_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":inner_product_op_py",
+        "//tensorflow_quantum/python:util",
+    ],
+)

tensorflow_quantum/core/ops/math_ops/__init__.py

@@ -0,0 +1,17 @@
+# 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 for tfq.core.ops.math_ops.*"""
+
+from tensorflow_quantum.core.ops.math_ops.inner_product_op import inner_product

tensorflow_quantum/core/ops/math_ops/inner_product_op.py

@@ -0,0 +1,55 @@
+# 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 python op gradient."""
+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 inner_product(programs, symbol_names, symbol_values, other_programs):
+    """Calculate the inner product between circuits.
+
+    Calculates out[i][j] = \langle \psi_{\text{programs[i]}} \\
+        (\text{symvol_values[i]}) | \psi_{\text{other_programs[j]}} \rangle
+
+    Note: `other_programs` must not contain any free symbols. These can resolved
+        beforehand with `tfq.resolve_parameters`.
+
+    Args:
+        programs: `tf.Tensor` of strings with shape [batch_size] containing
+            the string representations of the circuits
+        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`.
+        other_programs: `tf.Tensor` of strings with shape [batch_size, n_others]
+            containing the string representations of the circuits with which to
+            compute the overlap on `programs` with. Must not contain any free
+            symbols.
+    Returns:
+        `tf.Tensor` with shape [batch_size, n_others] where `out[i][j]` is equal
+            to the inner product of `programs[i]` with `symbol_values[i]`
+            resolved in and `other_programs[i][j]`.
+
+    """
+    return MATH_OP_MODULE.tfq_inner_product(programs, symbol_names,
+                                            tf.cast(symbol_values, tf.float32),
+                                            other_programs)