From 2ba9d7332cbd3b708b0f2b931f71168a61cc79b0 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Fri, 2 Apr 2021 22:45:45 +0900 Subject: [PATCH 01/12] Fix adjoint hessian bug --- tensorflow_quantum/core/ops/math_ops/BUILD | 12 + .../ops/math_ops/inner_product_grad_test.py | 10 +- .../math_ops/inner_product_hessian_test.py | 580 ++++++++++++ .../core/ops/math_ops/inner_product_op.py | 63 +- .../ops/math_ops/inner_product_op_test.py | 10 +- .../ops/math_ops/tfq_inner_product_hessian.cc | 839 ++++++++++++++++++ tensorflow_quantum/core/ops/parse_context.cc | 50 ++ tensorflow_quantum/core/ops/parse_context.h | 5 + tensorflow_quantum/core/src/BUILD | 17 + .../core/src/adj_hessian_util.cc | 611 +++++++++++++ .../core/src/adj_hessian_util.h | 124 +++ tensorflow_quantum/core/src/adj_util.cc | 2 - tensorflow_quantum/core/src/adj_util.h | 6 +- 13 files changed, 2310 insertions(+), 19 deletions(-) create mode 100644 tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py create mode 100644 tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc create mode 100644 tensorflow_quantum/core/src/adj_hessian_util.cc create mode 100644 tensorflow_quantum/core/src/adj_hessian_util.h diff --git a/tensorflow_quantum/core/ops/math_ops/BUILD b/tensorflow_quantum/core/ops/math_ops/BUILD index 537953e85..0f5a68ec0 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_inner_product_hessian.cc", ], copts = select({ ":windows": [ @@ -62,6 +63,7 @@ cc_binary( # cirq cc proto "//tensorflow_quantum/core/ops:parse_context", "//tensorflow_quantum/core/ops:tfq_simulate_utils", + "//tensorflow_quantum/core/src:adj_hessian_util", "//tensorflow_quantum/core/src:adj_util", "//tensorflow_quantum/core/src:circuit_parser_qsim", "//tensorflow_quantum/core/src:util_qsim", @@ -100,3 +102,13 @@ py_test( "//tensorflow_quantum/python:util", ], ) + +py_test( + name = "inner_product_hessian_test", + srcs = ["inner_product_hessian_test.py"], + python_version = "PY3", + deps = [ + ":inner_product_op_py", + "//tensorflow_quantum/python:util", + ], +) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_grad_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_grad_test.py index 6a4e8423e..6e5ce1bf1 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_grad_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_grad_test.py @@ -27,7 +27,7 @@ class InnerProductAdjGradTest(tf.test.TestCase, parameterized.TestCase): """Tests tfq_inner_product_grad.""" def test_inner_product_grad_inputs(self): - """Makes sure that inner_product_adj_grad fails on bad inputs.""" + """Makes sure that inner_product_grad fails on bad inputs.""" n_qubits = 5 batch_size = 5 n_other_programs = 3 @@ -232,7 +232,7 @@ def test_inner_product_grad_inputs(self): ]) def test_correctness_with_symbols(self, n_qubits, batch_size, inner_dim_size): - """Tests that inner_product works with symbols.""" + """Tests that inner_product_grad works with symbols.""" symbol_names = ['alpha', 'beta', 'gamma'] n_params = len(symbol_names) qubits = cirq.GridQubit.rect(1, n_qubits) @@ -242,7 +242,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size, other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] - for i in range(batch_size) + for _ in range(batch_size) ] symbol_values_array = np.array( @@ -312,7 +312,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size, ]) def test_correctness_without_symbols(self, n_qubits, batch_size, inner_dim_size): - """Tests that inner_product_adj_grad works without symbols.""" + """Tests that inner_product_grad works without symbols.""" qubits = cirq.GridQubit.rect(1, n_qubits) circuit_batch, _ = \ util.random_circuit_resolver_batch( @@ -320,7 +320,7 @@ def test_correctness_without_symbols(self, n_qubits, batch_size, other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] - for i in range(batch_size) + for _ in range(batch_size) ] programs = util.convert_to_tensor(circuit_batch) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py new file mode 100644 index 000000000..8d92b15df --- /dev/null +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -0,0 +1,580 @@ +# Copyright 2021 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 tfq_inner_product_hessian.""" +import copy +import time + +import numpy as np +from absl.testing import parameterized +import sympy +import tensorflow as tf +import cirq + +from tensorflow_quantum.core.ops.math_ops import inner_product_op +from tensorflow_quantum.python import util + +_INVERSE_SPEEDUP = 1 / 20.0 +_ATOL = 3e-3 +_SYMBOL_NAMES = [['alpha'], ['alpha', 'beta']] +_ONE_EIGEN_GATES = [ + cirq.XPowGate, + cirq.YPowGate, + cirq.ZPowGate, + cirq.HPowGate, + cirq.PhasedXPowGate, +] +_TWO_EIGEN_GATES = [ + cirq.XXPowGate, + cirq.YYPowGate, + cirq.ZZPowGate, + cirq.CZPowGate, + cirq.CNotPowGate, + cirq.SwapPowGate, + cirq.ISwapPowGate, + cirq.PhasedISwapPowGate, + cirq.FSimGate, +] + +_ATOL_FOR_COMPLEX_GATE = 1e-1 +_COMPLEX_GATES = [ + cirq.PhasedXPowGate, +] + + +def get_gate(gate, symbol_names, qubits): + symbols = sympy.symbols(symbol_names) + if len(symbols) == 1: + a, b = symbols * 2 + else: + a, b = symbols + if gate == cirq.PhasedXPowGate or gate == cirq.PhasedISwapPowGate: + return [gate(phase_exponent=0.1 * a, exponent=0.2*b).on(*qubits), + gate(phase_exponent=0.3 * b, exponent=0.4*a).on(*qubits)] + if gate == cirq.FSimGate: + return [gate(theta=0.1 * a, phi=0.2*b).on(*qubits), + gate(theta=0.3 * b, phi=0.4*a).on(*qubits)] + + return [gate(exponent=0.1 * a).on(*qubits), + gate(exponent=0.2 * b).on(*qubits)] + + +def get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, resolver): + new_resolver = copy.deepcopy(resolver) + if name_j == name_k: + new_resolver.param_dict[name_j] += (dx_j + dx_k) + else: + new_resolver.param_dict[name_j] += dx_j + new_resolver.param_dict[name_k] += dx_k + return cirq.resolve_parameters(circuit, new_resolver) + + +def get_finite_difference_hessian(circuit, name_j, name_k, resolver): + # dx came from _GRAD_EPS of core/src/adj_util.cc + dx = 5e-3 + inv_square_two_dx = np.asarray([1e4+0.j], dtype=np.complex64) + final_circuit_pp = get_shifted_resolved_circuit( + circuit, name_j, name_k, dx, dx, resolver) + final_circuit_mp = get_shifted_resolved_circuit( + circuit, name_j, name_k, -dx, dx, resolver) + final_circuit_pm = get_shifted_resolved_circuit( + circuit, name_j, name_k, dx, -dx, resolver) + final_circuit_mm = get_shifted_resolved_circuit( + circuit, name_j, name_k, -dx, -dx, + resolver) + final_wf_pp = inv_square_two_dx * cirq.final_state_vector(final_circuit_pp) + final_wf_mp = inv_square_two_dx * cirq.final_state_vector(final_circuit_mp) + final_wf_pm = inv_square_two_dx * cirq.final_state_vector(final_circuit_pm) + final_wf_mm = inv_square_two_dx * cirq.final_state_vector(final_circuit_mm) + # Performs central finite difference. + final_wf_grad = ((final_wf_pp + final_wf_mm) - (final_wf_pm + final_wf_mp)) + return final_wf_grad + + +class InnerProductAdjHessianTest(tf.test.TestCase, parameterized.TestCase): + """Tests tfq_inner_product_hessian.""" + # + # def test_inner_product_hessian_inputs(self): + # """Makes sure that inner_product_adj_hessian fails on bad inputs.""" + # n_qubits = 5 + # batch_size = 5 + # n_other_programs = 3 + # symbol_names = ['alpha'] + # qubits = cirq.GridQubit.rect(1, n_qubits) + # programs_coeffs = np.ones((batch_size,)) + # other_programs_coeffs = np.ones((batch_size, n_other_programs)) + # circuit_batch, resolver_batch = \ + # util.random_symbol_circuit_resolver_batch( + # qubits, symbol_names, batch_size) + # + # symbol_values_array = np.array( + # [[resolver[symbol] + # for symbol in symbol_names] + # for resolver in resolver_batch]) + # + # other_batch = [ + # util.random_circuit_resolver_batch(qubits, n_other_programs)[0] + # for _ in range(batch_size) + # ] + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'programs must be rank 1'): + # # Circuit tensor has too many dimensions. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor([circuit_batch]), + # symbol_names, symbol_values_array, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'symbol_names must be rank 1.'): + # # symbol_names tensor has too many dimensions. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), + # np.array([symbol_names]), symbol_values_array, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'symbol_values must be rank 2.'): + # # symbol_values_array tensor has too many dimensions. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # np.array([symbol_values_array]), + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'symbol_values must be rank 2.'): + # # symbol_values_array tensor has too few dimensions. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), + # symbol_names, symbol_values_array[0], + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'other_programs must be rank 2.'): + # # other_programs tensor has too few dimensions. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), + # symbol_names, symbol_values_array, + # util.convert_to_tensor(circuit_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'other_programs must be rank 2.'): + # # pauli_sums tensor has too many dimensions. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array, + # util.convert_to_tensor([[x] for x in other_batch]), + # programs_coeffs, other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'Unparseable proto'): + # # circuit tensor has the right type but invalid values. + # inner_product_op.inner_product_hessian( + # ['junk'] * batch_size, symbol_names, symbol_values_array, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'Could not find symbol in parameter map'): + # # symbol_names tensor has the right type but invalid values. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), + # ['junk'], symbol_values_array, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'not found in reference circuit'): + # # other_programs tensor has the right type but operates on + # # qubits that the reference ciruit doesn't have. + # new_qubits = [cirq.GridQubit(5, 5), cirq.GridQubit(9, 9)] + # new_circuits, _ = util.random_circuit_resolver_batch( + # new_qubits, batch_size) + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array, + # util.convert_to_tensor([[x] for x in new_circuits]), + # programs_coeffs, other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'not found in paired circuit'): + # # other_programs tensor has the right type but operates on + # # qubits that the reference ciruit doesn't have. + # new_qubits = cirq.GridQubit.rect(1, n_qubits - 1) + # new_circuits, _ = util.random_circuit_resolver_batch( + # new_qubits, batch_size) + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array, + # util.convert_to_tensor([[x] for x in new_circuits]), + # programs_coeffs, other_programs_coeffs) + # + # with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # # circuits tensor has the wrong type. + # inner_product_op.inner_product_hessian( + # [1.0] * batch_size, symbol_names, symbol_values_array, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # # symbol_names tensor has the wrong type. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), + # [0.1234], symbol_values_array, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.UnimplementedError, ''): + # # symbol_values tensor has the wrong type. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), + # symbol_names, [['junk']] * batch_size, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # # other_programs tensor has the wrong type. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array, [[1.0]] * batch_size, programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex(TypeError, 'missing'): + # # we are missing an argument. + # # pylint: disable=no-value-for-parameter + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array, programs_coeffs, other_programs_coeffs) + # # pylint: enable=no-value-for-parameter + # + # with self.assertRaisesRegex(TypeError, 'positional arguments'): + # # pylint: disable=too-many-function-args + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), + # symbol_names, symbol_values_array, + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs, []) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # expected_regex='do not match'): + # # batch programs has wrong batch size. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array, + # util.convert_to_tensor(other_batch[:int(batch_size * 0.5)]), + # programs_coeffs, other_programs_coeffs) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # expected_regex='do not match'): + # # batch programs has wrong batch size. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array[::int(batch_size * 0.5)], + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # + # with self.assertRaisesRegex( + # tf.errors.InvalidArgumentError, + # expected_regex='Found symbols in other_programs'): + # # other_programs has symbols. + # inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array, + # util.convert_to_tensor([[x] for x in circuit_batch]), + # programs_coeffs, other_programs_coeffs) + # + # res = inner_product_op.inner_product_hessian( + # util.convert_to_tensor(circuit_batch), symbol_names, + # symbol_values_array.astype(np.float64), + # util.convert_to_tensor(other_batch), programs_coeffs, + # other_programs_coeffs) + # self.assertDTypeEqual(res, np.complex64) + # + @parameterized.parameters([ + { + 'n_qubits': 5, + 'batch_size': 1, + 'inner_dim_size': 5 + }, + # { + # 'n_qubits': 5, + # 'batch_size': 10, + # 'inner_dim_size': 1 + # }, + # { + # 'n_qubits': 10, + # 'batch_size': 10, + # 'inner_dim_size': 2 + # }, + # { + # 'n_qubits': 5, + # 'batch_size': 10, + # 'inner_dim_size': 5 + # }, + ]) + def correctness_with_symbols(self, n_qubits, batch_size, + inner_dim_size): + """Tests that inner_product works with symbols.""" + symbol_names = ['alpha', 'beta', 'gamma'] + n_params = len(symbol_names) + qubits = cirq.GridQubit.rect(1, n_qubits) + circuit_batch, resolver_batch = \ + util.random_symbol_circuit_resolver_batch( + qubits, symbol_names, batch_size, n_moments=2) + print(circuit_batch) + + other_batch = [ + util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] + for _ in range(batch_size) + ] + + symbol_values_array = np.array( + [[resolver[symbol] + for symbol in symbol_names] + for resolver in resolver_batch]) + + programs = util.convert_to_tensor(circuit_batch) + other_programs = util.convert_to_tensor(other_batch) + symbol_names_tensor = tf.convert_to_tensor(symbol_names, + dtype=tf.dtypes.string) + symbol_values = tf.convert_to_tensor(symbol_values_array) + programs_coeffs = tf.cast(tf.random.normal((batch_size,)), tf.complex64) + other_programs_coeffs = tf.cast( + tf.random.normal((batch_size, inner_dim_size)), tf.complex64) + + t1 = time.time() + out = inner_product_op.inner_product_hessian( + programs, symbol_names_tensor, symbol_values, other_programs, + programs_coeffs, other_programs_coeffs) + t1 = time.time() - t1 + + t2 = time.time() + out_arr = np.zeros((batch_size, n_params, n_params), dtype=np.complex64) + for i, resolver in enumerate(resolver_batch): + weighted_internal_wf = None + for l, other in enumerate(other_batch[i]): + internal_wf = (other_programs_coeffs[i][l] * + cirq.final_state_vector(other)) + if l == 0: + weighted_internal_wf = internal_wf + else: + weighted_internal_wf += internal_wf + for j, name_j in enumerate(symbol_names): + for k, name_k in enumerate(symbol_names): + final_wf_grad = get_finite_difference_hessian( + circuit_batch[i], name_j, name_k, resolver) + out_arr[i][j][k] += ( + programs_coeffs[i] * + np.vdot(final_wf_grad, weighted_internal_wf)) + + # Elapsed time should be less than 5% of cirq version. + # (at least 20x speedup) + self.assertLess(t1, t2*_INVERSE_SPEEDUP) + self.assertAllClose(out, out_arr, atol=_ATOL) + + # + # @parameterized.parameters([ + # { + # 'n_qubits': 5, + # 'batch_size': 1, + # 'inner_dim_size': 5 + # }, + # { + # 'n_qubits': 5, + # 'batch_size': 10, + # 'inner_dim_size': 1 + # }, + # { + # 'n_qubits': 10, + # 'batch_size': 10, + # 'inner_dim_size': 2 + # }, + # { + # 'n_qubits': 5, + # 'batch_size': 10, + # 'inner_dim_size': 5 + # }, + # ]) + # def test_correctness_without_symbols(self, n_qubits, batch_size, + # inner_dim_size): + # """Tests that inner_product_adj_grad works without symbols.""" + # qubits = cirq.GridQubit.rect(1, n_qubits) + # circuit_batch, _ = \ + # util.random_circuit_resolver_batch( + # qubits, batch_size) + # + # other_batch = [ + # util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] + # for _ in range(batch_size) + # ] + # + # programs = util.convert_to_tensor(circuit_batch) + # other_programs = util.convert_to_tensor(other_batch) + # symbol_names = tf.convert_to_tensor([], dtype=tf.dtypes.string) + # symbol_values = tf.convert_to_tensor([[] for _ in range(batch_size)]) + # progams_coeffs = np.ones((batch_size,)) + # other_programs_coeffs = np.ones((batch_size, inner_dim_size)) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'symbols must be a positive integer'): + # inner_product_op.inner_product_hessian(programs, symbol_names, + # symbol_values, + # other_programs, + # progams_coeffs, + # other_programs_coeffs) + + # def test_correctness_empty(self): + # """Tests the inner product adj grad between two empty circuits.""" + # symbol_names = ['alpha', 'beta'] + # empty_cicuit = util.convert_to_tensor([cirq.Circuit()]) + # empty_symbols = tf.convert_to_tensor([], dtype=tf.dtypes.string) + # empty_values = tf.convert_to_tensor([[]]) + # other_program = util.convert_to_tensor([[cirq.Circuit()]]) + # program_coeffs = np.ones((1,)) + # other_program_coeffs = np.ones((1, 1)) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'symbols must be a positive integer'): + # inner_product_op.inner_product_hessian(empty_cicuit, empty_symbols, + # empty_values, other_program, + # program_coeffs, + # other_program_coeffs) + # + # empty_cicuit = util.convert_to_tensor([cirq.Circuit()]) + # symbol_names = tf.convert_to_tensor(symbol_names, + # dtype=tf.dtypes.string) + # symbol_values = tf.convert_to_tensor([[0.0 for _ in range(2)]]) + # other_program = util.convert_to_tensor([[cirq.Circuit()]]) + # + # out = inner_product_op.inner_product_hessian(empty_cicuit, symbol_names, + # symbol_values, + # other_program, + # program_coeffs, + # other_program_coeffs) + # expected = np.zeros((1, len(symbol_names)), dtype=np.complex64) + # self.assertAllClose(out, expected) + + # def test_correctness_no_circuit(self): + # """Test the inner product grad between no 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) + # other_program = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.string) + # empty_program_coeffs = tf.raw_ops.Empty(shape=(0,), dtype=tf.float32) + # empty_other_program_coeffs = tf.raw_ops.Empty(shape=(0, 0), + # dtype=tf.float32) + # + # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + # 'number of symbols must be a positive'): + # # When using `tf.gradients`, a user will never encounter this error + # # thanks to the `tf.cond` inside of the custom gradient. + # _ = inner_product_op.inner_product_hessian( + # empty_circuit, empty_symbols, empty_values, other_program, + # empty_program_coeffs, empty_other_program_coeffs) + + +class InnerProductHessianOnGates(tf.test.TestCase, parameterized.TestCase): + + @parameterized.parameters([ + { + 'gate': [cirq.XPowGate], + 'symbol_names': ['alpha','beta','gamma']# names + }]) # for gate in _ONE_EIGEN_GATES + _TWO_EIGEN_GATES for names in _SYMBOL_NAMES]) + def test_correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): + """Tests that inner_product works with symbols.""" + n_params = len(symbol_names) + qubits = cirq.GridQubit.rect(1, 5) # 2 if gate in _TWO_EIGEN_GATES else 1) + # circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] + circuit_batch = [cirq.Circuit([ + cirq.Moment( + (cirq.H**sympy.Mul(sympy.Float('0.96078327350981163', precision=53), sympy.Symbol('gamma'))).on(cirq.GridQubit(0, 0)), + (cirq.Y**sympy.Mul(sympy.Float('0.73193316007366105', precision=53), sympy.Symbol('alpha'))).on(cirq.GridQubit(0, 3)), + ), + cirq.Moment( + cirq.Y(cirq.GridQubit(0, 1)), + cirq.FSimGate(theta=0.123, phi=0.456).on(cirq.GridQubit(0, 4), cirq.GridQubit(0, 3)), + cirq.PhasedXPowGate(phase_exponent=0.123).on(cirq.GridQubit(0, 0)), + ), + cirq.Moment( + cirq.Y(cirq.GridQubit(0, 4)), + cirq.FSimGate(theta=0.123, phi=0.456).on(cirq.GridQubit(0, 2), cirq.GridQubit(0, 3)), + ), + cirq.Moment( + (cirq.H**sympy.Symbol('beta')).on(cirq.GridQubit(0, 0)), + ), + ])] + + resolver_batch = [cirq.ParamResolver({name: 0.123 for name in symbol_names})] + + symbol_values_array = np.array( + [[resolver[symbol] + for symbol in symbol_names] + for resolver in resolver_batch]) + other_batch = [ + [cirq.Circuit(cirq.H.on_each(*qubits))] + for _ in circuit_batch + ] + programs = util.convert_to_tensor(circuit_batch) + other_programs = util.convert_to_tensor(other_batch) + symbol_names_tensor = tf.convert_to_tensor(symbol_names, + dtype=tf.dtypes.string) + symbol_values = tf.convert_to_tensor(symbol_values_array) + programs_coeffs = tf.cast(tf.ones((1,)), tf.complex64) + other_programs_coeffs = tf.cast(tf.ones((1, 1)), tf.complex64) + # programs_coeffs = tf.cast(tf.random.normal((1,)), tf.complex64) + # other_programs_coeffs = tf.cast( + # tf.random.normal((1, 1)), tf.complex64) + + t1 = time.time() + out = inner_product_op.inner_product_hessian( + programs, symbol_names_tensor, symbol_values, other_programs, + programs_coeffs, other_programs_coeffs) + t1 = time.time() - t1 + + t2 = time.time() + out_arr = np.zeros((1, n_params, n_params), dtype=np.complex64) + for i, resolver in enumerate(resolver_batch): + weighted_internal_wf = None + for l, other in enumerate(other_batch[i]): + internal_wf = (other_programs_coeffs[i][l] * + cirq.final_state_vector(other)) + if l == 0: + weighted_internal_wf = internal_wf + else: + weighted_internal_wf += internal_wf + for j, name_j in enumerate(symbol_names): + out_arr[i][j][j] = inner_product_op._inner_product_grad( + programs, symbol_names_tensor, symbol_values, other_programs, + other_programs_coeffs)[i][j] + # for k, name_k in enumerate(symbol_names): + # final_wf_grad = get_finite_difference_hessian( + # circuit_batch[i], name_j, name_k, resolver) + # out_arr[i][j][k] += ( + # programs_coeffs[i] * + # np.vdot(final_wf_grad, weighted_internal_wf)) + + # Elapsed time should be less than 5% of cirq version. + # (at least 20x speedup) + self.assertLess(t1, t2*_INVERSE_SPEEDUP) + self.assertAllClose(out, out_arr, atol=_ATOL_FOR_COMPLEX_GATE if gate in _COMPLEX_GATES else _ATOL) + + +if __name__ == "__main__": + tf.test.main() diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py index 17675e978..0e2c8d961 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py @@ -20,6 +20,58 @@ MATH_OP_MODULE = load_module(os.path.join("math_ops", "_tfq_math_ops.so")) +def inner_product_hessian(programs, symbol_names, symbol_values, other_programs, + programs_coeffs, other_programs_coeffs): + """Calculate the adjoint Hessian of the inner product between circuits. + + Compute the gradients of the (potentially many) inner products between + the given circuits and the symbol free comparison circuits. + + Calculates out[i][j][k] = $\text{programs_coeffs[i]} \times \langle + \frac{\partial^2 \psi_{\text{programs[i]}}(\text{symbol_values[i]})} + {\partial \text{symbol_names[j]}\partial \text{symbol_names[k]}} | + \sum_l \text{other_programs_coeffs[l]}\times | + \psi_{\text{other_programs[l]}} \rangle$ + + + Note: `other_programs` must not contain any free symbols. These can + be resolved beforehand with `tfq.resolve_parameters`. + + Note: len(symbol_names) (=n_params) should be a positive integer. + + 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. + programs_coeffs: `tf.Tensor` of real numbers with shape [batch_size] + of weights on `programs`. + other_programs_coeffs: `tf.Tensor` of real numbers with shape + [batch_size, n_others] of weights on `other_programs`. + + Returns: + tf.Tensor` with shape [batch_size, n_symbols, n_symbols] where `out[i]` is + equal to the hessian of the inner product between programs[i] and all + other_programs[i] w.r.t. `symbol_names[j]` and `symbol_names[k]`. + `programs[i]` is resolved with `symbol_values[i]` and each + (other_)programs[i] is weighted by (other_)programs_coeffs[i]. + """ + return MATH_OP_MODULE.tfq_inner_product_hessian( + programs, symbol_names, tf.cast(symbol_values, tf.float32), + other_programs, tf.cast(programs_coeffs, tf.float32), + tf.cast(other_programs_coeffs, tf.float32)) + + def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, prev_grad): """Calculate the adjoint gradients of the inner product between circuits. @@ -27,9 +79,10 @@ def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, Compute the gradients of the (potentially many) inner products between the given circuits and the symbol free comparison circuits. - Calculates out[i][j][k] = $ \frac{\langle \psi_{\text{programs[i]}} \\ - (\text{symbol_values[i]})}{\partial \text{symbol_names[k]}} | \\ - \psi_{\text{other_programs[j]}} \rangle $ + Calculates out[i][j][k] = $\langle \frac{\partial \psi_{\text{programs[i]}} + (\text{symbol_values[i]})}{\partial \text{symbol_names[j]}} | \sum_k + \text{prev_grad[i][k]}\times | \psi_{\text{other_programs[k]}} + \rangle$ Note: `other_programs` must not contain any free symbols. These can @@ -40,7 +93,7 @@ def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, 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 + symbol_names: `tf.Tensor` of strings with shape [n_symbols], which is used to specify the order in which the values in `symbol_values` should be placed inside of the circuits in `programs`. @@ -52,7 +105,7 @@ def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, containing the string representations of the circuits with which to compute the overlap on `programs` with. Must not contain any free symbols. - prev_grad: `tf.Tensor` of real numbers with shape [batch_size, n_ops] + prev_grad: `tf.Tensor` of real numbers with shape [batch_size, n_others] backprop of values from downstream in the compute graph. Returns: diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_op_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_op_test.py index ee9fd08b4..95bc1e023 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_op_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_op_test.py @@ -43,7 +43,7 @@ def test_inner_product_inputs(self): other_batch = [ util.random_circuit_resolver_batch(qubits, 3)[0] - for i in range(batch_size) + for _ in range(batch_size) ] with self.assertRaisesRegex(tf.errors.InvalidArgumentError, @@ -239,7 +239,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size, other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] - for i in range(batch_size) + for _ in range(batch_size) ] symbol_values_array = np.array( @@ -299,7 +299,7 @@ def test_correctness_without_symbols(self, n_qubits, batch_size, other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] - for i in range(batch_size) + for _ in range(batch_size) ] programs = util.convert_to_tensor(circuit_batch) @@ -368,7 +368,7 @@ def test_tf_gradient_correctness_with_symbols(self, n_qubits, batch_size, other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] - for i in range(batch_size) + for _ in range(batch_size) ] symbol_values_array = np.array( @@ -434,7 +434,7 @@ def test_tf_gradient_correctness_without_symbols(self, n_qubits, batch_size, other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] - for i in range(batch_size) + for _ in range(batch_size) ] programs = util.convert_to_tensor(circuit_batch) diff --git a/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc new file mode 100644 index 000000000..427528258 --- /dev/null +++ b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc @@ -0,0 +1,839 @@ +/* Copyright 2021 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 +#include +#include + +#include "../qsim/lib/circuit.h" +#include "../qsim/lib/gate_appl.h" +#include "../qsim/lib/gates_cirq.h" +#include "../qsim/lib/seqfor.h" +#include "../qsim/lib/simmux.h" +#include "cirq/google/api/v2/program.pb.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/src/adj_hessian_util.h" +#include "tensorflow_quantum/core/src/adj_util.h" +#include "tensorflow_quantum/core/src/util_qsim.h" + +namespace tfq { + +using ::cirq::google::api::v2::Program; +using ::tensorflow::Status; +using ::tfq::proto::PauliSum; + +typedef qsim::Cirq::GateCirq QsimGate; +typedef qsim::Circuit QsimCircuit; +typedef std::vector> QsimFusedCircuit; + +class TfqInnerProductHessianOp : public tensorflow::OpKernel { + public: + explicit TfqInnerProductHessianOp(tensorflow::OpKernelConstruction* context) + : OpKernel(context) {} + + 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 == 6, + tensorflow::errors::InvalidArgument(absl::StrCat( + "Expected 6 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_internal_size = context->input(3).dim_size(1); + const int output_dim_symbol_size = context->input(1).dim_size(0); + OP_REQUIRES(context, output_dim_symbol_size > 0, + tensorflow::errors::InvalidArgument(absl::StrCat( + "The number of symbols must be a positive integer, got ", + output_dim_symbol_size, " symbols."))); + tensorflow::TensorShape output_shape; + output_shape.AddDim(output_dim_batch_size); + output_shape.AddDim(output_dim_symbol_size); + output_shape.AddDim(output_dim_symbol_size); + + tensorflow::Tensor* output = nullptr; + OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); + auto output_tensor = output->tensor, 3>(); + output_tensor.setZero(); + + // Parse program protos. + std::vector programs; + std::vector num_qubits; + std::vector> other_programs; + OP_REQUIRES_OK(context, + GetProgramsAndNumQubits(context, &programs, &num_qubits, + &other_programs)); + + std::vector 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."))); + OP_REQUIRES(context, output_dim_symbol_size == maps[0].size(), + tensorflow::errors::InvalidArgument(absl::StrCat( + "Number of symbols and symbol maps do not match. Got ", + output_dim_symbol_size, " symbols and ", maps[0].size(), + " symbol values."))); + + // Construct qsim circuits for programs. + std::vector qsim_circuits(programs.size(), QsimCircuit()); + std::vector fused_circuits(programs.size(), + QsimFusedCircuit({})); + + // track metadata. + std::vector> gate_meta( + programs.size(), std::vector({})); + + // Construct qsim circuits. + std::vector>>> + partial_fused_grad_circuits( + programs.size(), + std::vector>>({})); + std::vector>>> + partial_fused_hess_circuits( + programs.size(), + std::vector>>({})); + + // track gradients + // For 1st order gradient gates + std::vector> gradient_gates( + programs.size(), std::vector({})); + + // For 2nd order gradient gates + std::vector> hessian_gates( + programs.size(), std::vector({})); + + 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], &gate_meta[i]); + NESTED_FN_STATUS_SYNC(parse_status, local, p_lock); + + CreateGradientCircuit(qsim_circuits[i], gate_meta[i], + &partial_fused_grad_circuits[i], + &gradient_gates[i]); + CreateHessianCircuit(qsim_circuits[i], gate_meta[i], + &partial_fused_hess_circuits[i], + &hessian_gates[i]); + } + }; + + 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); + + // Construct qsim circuits for other_programs. + std::vector> other_qsim_circuits( + output_dim_batch_size, + std::vector(output_dim_internal_size, QsimCircuit())); + std::vector> other_fused_circuits( + output_dim_batch_size, + std::vector(output_dim_internal_size, + QsimFusedCircuit({}))); + + auto construct_f2 = [&](int start, int end) { + for (int i = start; i < end; i++) { + int ii = i / output_dim_internal_size; + int jj = i % output_dim_internal_size; + Status status = QsimCircuitFromProgram( + other_programs[ii][jj], {}, num_qubits[ii], + &other_qsim_circuits[ii][jj], &other_fused_circuits[ii][jj]); + NESTED_FN_STATUS_SYNC(parse_status, status, p_lock); + } + }; + + context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor( + output_dim_batch_size * output_dim_internal_size, num_cycles, + construct_f2); + if (!parse_status.ok()) { + OP_REQUIRES_OK(context, + tensorflow::errors::InvalidArgument(absl::StrCat( + "Found symbols in other_programs.", + "No symbols are allowed in these circuits."))); + } + + // Get programs coefficients. + std::vector programs_coeffs; + std::vector> other_programs_coeffs; + OP_REQUIRES_OK(context, GetProgramsCoefficients(context, &programs_coeffs, + &other_programs_coeffs)); + + int max_num_qubits = 0; + for (const int num : num_qubits) { + max_num_qubits = std::max(max_num_qubits, num); + } + + output_tensor.setZero(); + + // Cross reference with standard google cloud compute instances + // Memory ~= 2 * num_threads * (2 * 64 * 2 ** num_qubits in circuits) + // e2s2 = 2 CPU, 8GB -> Can safely do 22 since Memory = 4GB + // e2s4 = 4 CPU, 16GB -> Can safely do 22 since Memory = 8GB + // ... + if (true || max_num_qubits >= 22 || output_dim_batch_size == 1) { + ComputeLarge(num_qubits, maps, qsim_circuits, fused_circuits, + partial_fused_grad_circuits, partial_fused_hess_circuits, + gradient_gates, hessian_gates, other_fused_circuits, + programs_coeffs, other_programs_coeffs, context, + &output_tensor); + } else { + ComputeSmall(num_qubits, max_num_qubits, maps, qsim_circuits, + fused_circuits, partial_fused_grad_circuits, + partial_fused_hess_circuits, gradient_gates, hessian_gates, + other_fused_circuits, programs_coeffs, other_programs_coeffs, + context, &output_tensor); + } + } + + private: + void ComputeLarge( + const std::vector& num_qubits, const std::vector& maps, + const std::vector& qsim_circuits, + const std::vector& fused_circuits, + const std::vector>>>& + partial_fused_grad_circuits, + const std::vector>>>& + partial_fused_hess_circuits, + const std::vector>& gradient_gates, + const std::vector>& hessian_gates, + const std::vector>& other_fused_circuits, + const std::vector& programs_coeffs, + const std::vector>& other_programs_coeffs, + tensorflow::OpKernelContext* context, + tensorflow::TTypes, 3>::Tensor* output_tensor) { + // Instantiate qsim objects. + const auto tfq_for = tfq::QsimFor(context); + using Simulator = qsim::Simulator; + using StateSpace = Simulator::StateSpace; + + // Begin simulation. + int largest_nq = 1; + Simulator sim = Simulator(tfq_for); + StateSpace ss = StateSpace(tfq_for); + auto sv = ss.Create(largest_nq); + auto scratch = ss.Create(largest_nq); + auto scratch2 = ss.Create(largest_nq); + auto scratch3 = ss.Create(largest_nq); + auto scratch4 = ss.Create(largest_nq); + + // Simulate programs one by one. Parallelizing over state vectors + // we no longer parallelize over circuits. Each time we encounter a + // a larger circuit we will grow the Statevector as necessary. + for (std::vector>>::size_type i = 0; + i < fused_circuits.size(); i++) { + int nq = num_qubits[i]; + if (nq > largest_nq) { + // need to switch to larger statespace. + largest_nq = nq; + sv = ss.Create(largest_nq); + scratch = ss.Create(largest_nq); + scratch2 = ss.Create(largest_nq); + scratch3 = ss.Create(largest_nq); + scratch4 = ss.Create(largest_nq); + } + ss.SetStateZero(sv); + for (std::vector>::size_type j = 0; + j < fused_circuits[i].size(); j++) { + qsim::ApplyFusedGate(sim, fused_circuits[i][j], sv); + } + + auto status = AccumulateFusedCircuits(other_programs_coeffs[i], + other_fused_circuits[i], sim, ss, + scratch2, scratch); + + // now sv is |psi> + // scratch contains sum_j other_programs_coeffs[i][j]*|phi[i][j]> + // Start adjoint differentiation on a single gate + std::cout << ">>> Start a single gate hessian" << std::endl; + for (int l = partial_fused_hess_circuits[i].size() - 1; l >= 0; l--) { + std::cout << ">>>>>> " << l << "th partial fused circuit is applied" << std::endl; + for (int k = partial_fused_hess_circuits[i][l].size() - 1; k >= 0; + k--) { + ApplyFusedGateDagger(sim, partial_fused_hess_circuits[i][l][k], sv); + ApplyFusedGateDagger(sim, partial_fused_hess_circuits[i][l][k], + scratch); + } + if (l == 0) { + // last layer will have no parametrized gates so can break. + break; + } + + // Hit a parameterized gate. + auto cur_gate = qsim_circuits[i].gates[hessian_gates[i][l - 1].index]; + ApplyGateDagger(sim, cur_gate, sv); + + // if applicable compute control qubit mask and control value bits. + uint64_t mask = 0; + uint64_t cbits = 0; + for (std::vector::size_type k = 0; + k < cur_gate.controlled_by.size(); k++) { + uint64_t control_loc = cur_gate.controlled_by[k]; + mask |= uint64_t{1} << control_loc; + cbits |= ((cur_gate.cmask >> k) & 1) << control_loc; + } + + for (std::vector::size_type k = 0; + k < hessian_gates[i][l - 1].grad_gates.size(); k++) { + // Copy sv onto scratch2 in anticipation of non-unitary "gradient + // gate". + ss.Copy(sv, scratch2); + if (!cur_gate.controlled_by.empty()) { + // Gradient of controlled gates puts zeros on diagonal which is + // the same as collapsing the state and then applying the + // non-controlled version of the gradient gate. + ss.BulkSetAmpl(scratch2, mask, cbits, 0, 0, true); + } + std::cout << ">>>>>>... " << k << "th gradient gate is applied" << std::endl; + qsim::ApplyGate(sim, hessian_gates[i][l - 1].grad_gates[k], scratch2); + + auto ptr = scratch2.get(); + auto ptr_size = 2 << scratch2.num_qubits(); + std::cout << "Statevector" << std::endl; + for (int i = 0; i < ptr_size; i++) { + std::cout << ptr[i] << ","; + } + std::cout << std::endl; + + ptr = scratch.get(); + ptr_size = 2 << scratch.num_qubits(); + std::cout << "Other Statevector" << std::endl; + for (int i = 0; i < ptr_size; i++) { + std::cout << ptr[i] << ","; + } + std::cout << std::endl; + // don't need not-found check since this is done upstream already. + auto symbol = hessian_gates[i][l - 1].params[k]; + std::cout << ">>>>>>... " << k << "th symbol = " << symbol << std::endl; + double coeff = static_cast(programs_coeffs[i]); + if (symbol == kUsePrevTwoSymbols) { + // Apply second-order finite difference w.r.t. two symbols + // That is, CrossTerm w.r.t. two symbols in one gate. + auto symbol1 = hessian_gates[i][l - 1].params[k - 2]; + auto symbol2 = hessian_gates[i][l - 1].params[k - 1]; + std::cout << ">>>>>>... CrossTerm : two symbols are " << symbol1 << " and " << symbol2 << std::endl; + auto it = maps[i].find(symbol1); + const int loc1 = it->second.first; + it = maps[i].find(symbol2); + const int loc2 = it->second.first; + // Apply finite differencing for adjoint gradients. + // Finite differencing enables applying multiple `gradient_gate` + // of a symbol at the same circuit. For analytic methods like + // parameter-shift we need to apply a single `gradient_gate` + // per a symbol. + std::complex result = ss.InnerProduct(scratch2, scratch); + auto val = ( + std::complex(static_cast(coeff * result.real()), + static_cast(coeff * result.imag()))); + // Because Hessian is symmetric. + std::cout << ">>>>>>... value" << val << std::endl; + (*output_tensor)(i, loc1, loc2) += val; + std::cout << ">>>>>>... (*output_tensor)(i, loc1, loc2)" << (*output_tensor)(i, loc1, loc2) << std::endl; + (*output_tensor)(i, loc2, loc1) += val; + std::cout << ">>>>>>... (*output_tensor)(i, loc2, loc1)" << (*output_tensor)(i, loc2, loc1) << std::endl; + } else { + std::cout << ">>>>>>... One gate " << symbol << std::endl; + // Apply second-order finite difference w.r.t. one symbol + const auto it = maps[i].find(symbol); + const int loc = it->second.first; + // Apply finite differencing for adjoint gradients. + // Finite differencing enables applying multiple `gradient_gate` + // of a symbol at the same circuit. For analytic methods like + // parameter-shift we need to apply a single `gradient_gate` + // per a symbol. + std::complex result = ss.InnerProduct(scratch2, scratch); + auto val = + ( + std::complex(static_cast(coeff * result.real()), + static_cast(coeff * result.imag()))); + std::cout << ">>>>>>... value" << val << std::endl; + (*output_tensor)(i, loc, loc) += val; + } + } + ApplyGateDagger(sim, cur_gate, scratch); + } + + // Re-initialize statevectors to save memory. + ss.SetStateZero(sv); + for (std::vector>::size_type j = 0; + j < fused_circuits[i].size(); j++) { + qsim::ApplyFusedGate(sim, fused_circuits[i][j], sv); + } + + status = AccumulateFusedCircuits(other_programs_coeffs[i], + other_fused_circuits[i], sim, ss, + scratch2, scratch); + // now sv is |psi> + // scratch contains sum_j other_programs_coeffs[i][j]*|phi[i][j]> + // Start adjoint differentiation on two gates + // m is the index for the first gate + std::cout << ">>> Start two gates hessian" << std::endl; + for (int m = partial_fused_grad_circuits[i].size() - 1; m >= 1; m--) { + std::cout << ">>>>>>(1) m=" << m << "th partial fused circuit is applied" << std::endl; + for (int k = partial_fused_grad_circuits[i][m].size() - 1; k >= 0; + k--) { + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[i][m][k], sv); + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[i][m][k], + scratch); + } + auto cur_gate_m = + qsim_circuits[i].gates[gradient_gates[i][m - 1].index]; + ApplyGateDagger(sim, cur_gate_m, sv); + + // if applicable compute control qubit mask and control value bits. + uint64_t mask_m = 0; + uint64_t cbits_m = 0; + for (std::vector::size_type k = 0; + k < cur_gate_m.controlled_by.size(); k++) { + uint64_t control_loc = cur_gate_m.controlled_by[k]; + mask_m |= uint64_t{1} << control_loc; + cbits_m |= ((cur_gate_m.cmask >> k) & 1) << control_loc; + } + for (std::vector::size_type p = 0; + p < gradient_gates[i][m - 1].grad_gates.size(); p++) { + // Copy sv onto scratch2 in anticipation of the first non-unitary + // "gradient gate". + ss.Copy(sv, scratch2); + if (!cur_gate_m.controlled_by.empty()) { + // Gradient of controlled gates puts zeros on diagonal which is + // the same as collapsing the state and then applying the + // non-controlled version of the gradient gate. + ss.BulkSetAmpl(scratch2, mask_m, cbits_m, 0, 0, true); + } + std::cout << ">>>>>>(1)... p=" << p << "th gradient gate is applied" << std::endl; + qsim::ApplyGate(sim, gradient_gates[i][m - 1].grad_gates[p], + scratch2); + + // don't need not-found check since this is done upstream already. + const auto it = maps[i].find(gradient_gates[i][m - 1].params[p]); + std::cout << ">>>>>>(1)... p=" << p << "th symbol = " << gradient_gates[i][m - 1].params[p] << std::endl; + const int loc_m = it->second.first; + + // scratch2 is now (d/dsymbol[p])|psi> + // Copy scratch onto scratch4. + ss.Copy(scratch, scratch4); + // ApplyGateDagger(sim, cur_gate_m, scratch4); + // n is the index for the second gate + for (int n = m - 1; n >= 0; n--) { + std::cout << ">>>>>>---(2) " << n << "th partial fused circuit is applied" << std::endl; + for (int k = partial_fused_grad_circuits[i][n].size() - 1; k >= 0; + k--) { + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[i][n][k], + scratch2); + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[i][n][k], + scratch4); + } + if (n == 0) { + std::cout << ">>>>>>---(2) no just quit." << std::endl; + // last layer will have no parametrized gates so can break. + break; + } + + // Hit a parameterized gate. + auto cur_gate_n = + qsim_circuits[i].gates[gradient_gates[i][n - 1].index]; + ApplyGateDagger(sim, cur_gate_n, scratch2); + + // if applicable compute control qubit mask and control value bits. + uint64_t mask_n = 0; + uint64_t cbits_n = 0; + for (std::vector::size_type k = 0; + k < cur_gate_n.controlled_by.size(); k++) { + uint64_t control_loc = cur_gate_n.controlled_by[k]; + mask_n |= uint64_t{1} << control_loc; + cbits_n |= ((cur_gate_n.cmask >> k) & 1) << control_loc; + } + + for (std::vector::size_type q = 0; + q < gradient_gates[i][n - 1].grad_gates.size(); q++) { + // Copy scratch2 onto scratch3 in anticipation of the second + // non-unitary "gradient gate". + ss.Copy(scratch2, scratch3); + if (!cur_gate_n.controlled_by.empty()) { + // Gradient of controlled gates puts zeros on diagonal which is + // the same as collapsing the state and then applying the + // non-controlled version of the gradient gate. + ss.BulkSetAmpl(scratch3, mask_n, cbits_n, 0, 0, true); + } + + std::cout << ">>>>>>---(2)... q=" << q << "th gradient gate is applied" << std::endl; + qsim::ApplyGate(sim, gradient_gates[i][n - 1].grad_gates[q], + scratch3); + + // don't need not-found check since this is done upstream already. + const auto it = maps[i].find(gradient_gates[i][n - 1].params[q]); + std::cout << ">>>>>>---(2)... q=" << q << "th symbol = " << gradient_gates[i][n - 1].params[q] << std::endl; + const int loc_n = it->second.first; + // Apply finite differencing for adjoint gradients. + // Finite differencing enables applying multiple `gradient_gate` + // of a symbol at the same circuit. For analytic methods like + // parameter-shift we need to apply a single `gradient_gate` + // per a symbol. + double coeff = static_cast(programs_coeffs[i]); + std::complex result = ss.InnerProduct(scratch3, scratch4); + auto val = + ( + std::complex(static_cast(coeff * result.real()), + static_cast(coeff * result.imag()))); + std::cout << ">>>>>>>>>... value" << val << std::endl; + (*output_tensor)(i, loc_m, loc_n) += val; + (*output_tensor)(i, loc_n, loc_m) += val; + } + ApplyGateDagger(sim, cur_gate_n, scratch4); + } + } + ApplyGateDagger(sim, cur_gate_m, scratch); + } + } + } + + void ComputeSmall( + const std::vector& num_qubits, const int max_num_qubits, + const std::vector& maps, + const std::vector& qsim_circuits, + const std::vector& fused_circuits, + const std::vector>>>& + partial_fused_grad_circuits, + const std::vector>>>& + partial_fused_hess_circuits, + const std::vector>& gradient_gates, + const std::vector>& hessian_gates, + const std::vector>& other_fused_circuits, + const std::vector& programs_coeffs, + const std::vector>& other_programs_coeffs, + tensorflow::OpKernelContext* context, + tensorflow::TTypes, 3>::Tensor* output_tensor) { + const auto tfq_for = qsim::SequentialFor(1); + using Simulator = qsim::Simulator; + using StateSpace = Simulator::StateSpace; + + const int output_dim_internal_size = other_fused_circuits[0].size(); + + auto DoWork = [&](int start, int end) { + int old_batch_index = -2; + int cur_batch_index = -1; + int largest_nq = 1; + int cur_internal_index; + + Simulator sim = Simulator(tfq_for); + StateSpace ss = StateSpace(tfq_for); + auto sv = ss.Create(largest_nq); + auto sv_adj = ss.Create(largest_nq); + auto scratch = ss.Create(largest_nq); + auto scratch2 = ss.Create(largest_nq); + auto scratch3 = ss.Create(largest_nq); + auto scratch4 = ss.Create(largest_nq); + for (int i = start; i < end; i++) { + cur_batch_index = i / output_dim_internal_size; + cur_internal_index = i % output_dim_internal_size; + + const int nq = num_qubits[cur_batch_index]; + + 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); + sv_adj = ss.Create(largest_nq); + scratch = ss.Create(largest_nq); + scratch2 = ss.Create(largest_nq); + } + ss.SetStateZero(sv); + for (std::vector>::size_type j = 0; + j < fused_circuits[cur_batch_index].size(); j++) { + qsim::ApplyFusedGate(sim, fused_circuits[cur_batch_index][j], sv); + } + } + + ss.SetStateZero(scratch); + for (std::vector>::size_type k = 0; + k < + other_fused_circuits[cur_batch_index][cur_internal_index].size(); + k++) { + qsim::ApplyFusedGate( + sim, other_fused_circuits[cur_batch_index][cur_internal_index][k], + scratch); + } + + // now sv is |psi> + // scratch contains |phi> + // Start adjoint differentiation on a single gate + for (int l = partial_fused_hess_circuits[cur_batch_index].size() - 1; l >= 0; l--) { + for (int k = partial_fused_hess_circuits[cur_batch_index][l].size() - 1; k >= 0; + k--) { + ApplyFusedGateDagger(sim, partial_fused_hess_circuits[cur_batch_index][l][k], sv); + ApplyFusedGateDagger(sim, partial_fused_hess_circuits[cur_batch_index][l][k], + scratch); + } + if (l == 0) { + // last layer will have no parametrized gates so can break. + break; + } + + // Hit a parameterized gate. + auto cur_gate = qsim_circuits[cur_batch_index].gates[hessian_gates[cur_batch_index][l - 1].index]; + ApplyGateDagger(sim, cur_gate, sv); + + // if applicable compute control qubit mask and control value bits. + uint64_t mask = 0; + uint64_t cbits = 0; + for (std::vector::size_type k = 0; + k < cur_gate.controlled_by.size(); k++) { + uint64_t control_loc = cur_gate.controlled_by[k]; + mask |= uint64_t{1} << control_loc; + cbits |= ((cur_gate.cmask >> k) & 1) << control_loc; + } + + for (std::vector::size_type k = 0; + k < hessian_gates[cur_batch_index][l - 1].grad_gates.size(); k++) { + // Copy sv onto scratch2 in anticipation of non-unitary "gradient + // gate". + ss.Copy(sv, scratch2); + if (!cur_gate.controlled_by.empty()) { + // Gradient of controlled gates puts zeros on diagonal which is + // the same as collapsing the state and then applying the + // non-controlled version of the gradient gate. + ss.BulkSetAmpl(scratch2, mask, cbits, 0, 0, true); + } + qsim::ApplyGate(sim, hessian_gates[cur_batch_index][l - 1].grad_gates[k], scratch2); + + // don't need not-found check since this is done upstream already. + auto symbol = hessian_gates[cur_batch_index][l - 1].params[k]; + if (symbol == kUsePrevTwoSymbols) { + // Apply second-order finite difference w.r.t. two symbols + // That is, CrossTerm w.r.t. two symbols in one gate. + auto symbol1 = hessian_gates[cur_batch_index][l - 1].params[k - 2]; + auto symbol2 = hessian_gates[cur_batch_index][l - 1].params[k - 1]; + auto it = maps[cur_batch_index].find(symbol1); + const int loc1 = it->second.first; + it = maps[cur_batch_index].find(symbol2); + const int loc2 = it->second.first; + // Apply finite differencing for adjoint gradients. + // Finite differencing enables applying multiple `gradient_gate` + // of a symbol at the same circuit. For analytic methods like + // parameter-shift we need to apply a single `gradient_gate` + // per a symbol. + std::complex result = ss.InnerProduct(scratch2, scratch); + auto val = (programs_coeffs[cur_batch_index] * other_programs_coeffs[cur_batch_index][cur_internal_index] * + std::complex(static_cast(result.real()), + static_cast(result.imag()))); + // Because Hessian is symmetric. + (*output_tensor)(cur_batch_index, loc1, loc2) += val; + (*output_tensor)(cur_batch_index, loc2, loc1) += val; + } else { + // Apply second-order finite difference w.r.t. one symbol + const auto it = maps[cur_batch_index].find(symbol); + const int loc = it->second.first; + // Apply finite differencing for adjoint gradients. + // Finite differencing enables applying multiple `gradient_gate` + // of a symbol at the same circuit. For analytic methods like + // parameter-shift we need to apply a single `gradient_gate` + // per a symbol. + std::complex result = ss.InnerProduct(scratch2, scratch); + (*output_tensor)(cur_batch_index, loc, loc) += + (programs_coeffs[cur_batch_index] * other_programs_coeffs[cur_batch_index][cur_internal_index] * + std::complex(static_cast(result.real()), + static_cast(result.imag()))); + } + } + ApplyGateDagger(sim, cur_gate, scratch); + } + + // Re-initialize statevectors to save memory. + ss.SetStateZero(sv); + for (std::vector>::size_type j = 0; + j < fused_circuits[cur_batch_index].size(); j++) { + qsim::ApplyFusedGate(sim, fused_circuits[cur_batch_index][j], sv); + } + + // now sv is |psi> + // scratch contains |phi> + // Start adjoint differentiation on two gates + // m is the index for the first gate + for (int m = partial_fused_grad_circuits[cur_batch_index].size() - 1; m >= 1; m--) { + for (int k = partial_fused_grad_circuits[cur_batch_index][m].size() - 1; k >= 0; + k--) { + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][m][k], sv); + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][m][k], + scratch); + } + auto cur_gate_m = + qsim_circuits[cur_batch_index].gates[gradient_gates[cur_batch_index][m - 1].index]; + ApplyGateDagger(sim, cur_gate_m, sv); + + // if applicable compute control qubit mask and control value bits. + uint64_t mask_m = 0; + uint64_t cbits_m = 0; + for (std::vector::size_type k = 0; + k < cur_gate_m.controlled_by.size(); k++) { + uint64_t control_loc = cur_gate_m.controlled_by[k]; + mask_m |= uint64_t{1} << control_loc; + cbits_m |= ((cur_gate_m.cmask >> k) & 1) << control_loc; + } + for (std::vector::size_type p = 0; + p < gradient_gates[cur_batch_index][m - 1].grad_gates.size(); p++) { + // Copy sv onto scratch2 in anticipation of the first non-unitary + // "gradient gate". + ss.Copy(sv, scratch2); + if (!cur_gate_m.controlled_by.empty()) { + // Gradient of controlled gates puts zeros on diagonal which is + // the same as collapsing the state and then applying the + // non-controlled version of the gradient gate. + ss.BulkSetAmpl(scratch2, mask_m, cbits_m, 0, 0, true); + } + qsim::ApplyGate(sim, gradient_gates[cur_batch_index][m - 1].grad_gates[p], + scratch2); + + // don't need not-found check since this is done upstream already. + const auto it = maps[cur_batch_index].find(gradient_gates[cur_batch_index][m - 1].params[p]); + const int loc_m = it->second.first; + + // Copy scratch onto scratch4. + ss.Copy(scratch, scratch4); + // n is the index for the second gate + for (int n = m - 1; n >= 0; n--) { + for (int k = partial_fused_grad_circuits[cur_batch_index][n].size() - 1; k >= 0; + k--) { + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][n][k], + scratch2); + ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][n][k], + scratch4); + } + if (n == 0) { + // last layer will have no parametrized gates so can break. + break; + } + + // Hit a parameterized gate. + auto cur_gate_n = + qsim_circuits[cur_batch_index].gates[gradient_gates[cur_batch_index][n - 1].index]; + ApplyGateDagger(sim, cur_gate_n, scratch2); + + // if applicable compute control qubit mask and control value bits. + uint64_t mask_n = 0; + uint64_t cbits_n = 0; + for (std::vector::size_type k = 0; + k < cur_gate_n.controlled_by.size(); k++) { + uint64_t control_loc = cur_gate_n.controlled_by[k]; + mask_n |= uint64_t{1} << control_loc; + cbits_n |= ((cur_gate_n.cmask >> k) & 1) << control_loc; + } + + for (std::vector::size_type q = 0; + q < gradient_gates[cur_batch_index][n - 1].grad_gates.size(); q++) { + // Copy scratch2 onto scratch3 in anticipation of the second + // non-unitary "gradient gate". + ss.Copy(scratch2, scratch3); + if (!cur_gate_n.controlled_by.empty()) { + // Gradient of controlled gates puts zeros on diagonal which is + // the same as collapsing the state and then applying the + // non-controlled version of the gradient gate. + ss.BulkSetAmpl(scratch3, mask_n, cbits_n, 0, 0, true); + } + qsim::ApplyGate(sim, gradient_gates[cur_batch_index][n - 1].grad_gates[q], + scratch3); + + // don't need not-found check since this is done upstream already. + const auto it = maps[cur_batch_index].find(gradient_gates[cur_batch_index][n - 1].params[q]); + const int loc_n = it->second.first; + // Apply finite differencing for adjoint gradients. + // Finite differencing enables applying multiple `gradient_gate` + // of a symbol at the same circuit. For analytic methods like + // parameter-shift we need to apply a single `gradient_gate` + // per a symbol. + std::complex result = ss.InnerProduct(scratch3, scratch4); + auto val = + (programs_coeffs[cur_batch_index] * other_programs_coeffs[cur_batch_index][cur_internal_index] * + std::complex(static_cast(result.real()), + static_cast(result.imag()))); + (*output_tensor)(cur_batch_index, loc_m, loc_n) += val; + if (loc_m != loc_n) { + (*output_tensor)(cur_batch_index, loc_n, loc_m) += val; + } + } + ApplyGateDagger(sim, cur_gate_n, scratch4); + } + } + ApplyGateDagger(sim, cur_gate_m, scratch); + } + old_batch_index = cur_batch_index; + } + }; + + const int64_t num_cycles = + 200 * (int64_t(1) << static_cast(max_num_qubits)); + context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor( + fused_circuits.size() * output_dim_internal_size, num_cycles, DoWork); + } +}; + +REGISTER_KERNEL_BUILDER( + Name("TfqInnerProductHessian").Device(tensorflow::DEVICE_CPU), + TfqInnerProductHessianOp); + +REGISTER_OP("TfqInnerProductHessian") + .Input("programs: string") + .Input("symbol_names: string") + .Input("symbol_values: float") + .Input("other_programs: string") + .Input("programs_coeffs: float") + .Input("other_programs_coeffs: float") + .Output("inner_products_hessian: complex64") + .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 other_programs_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 2, &other_programs_shape)); + + tensorflow::shape_inference::ShapeHandle programs_coeffs_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 1, &programs_coeffs_shape)); + + tensorflow::shape_inference::ShapeHandle other_programs_coeffs_shape; + TF_RETURN_IF_ERROR( + c->WithRank(c->input(5), 2, &other_programs_coeffs_shape)); + + tensorflow::shape_inference::DimensionHandle output_rows = + c->Dim(programs_shape, 0); + tensorflow::shape_inference::DimensionHandle output_cols = + c->Dim(symbol_names_shape, 0); + std::vector dims = { + output_rows, output_cols, output_cols}; + c->set_output(0, c->MakeShape(dims)); + + 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 3caa4f3ed..55eb23003 100644 --- a/tensorflow_quantum/core/ops/parse_context.cc +++ b/tensorflow_quantum/core/ops/parse_context.cc @@ -425,4 +425,54 @@ tensorflow::Status GetPrevGrads( return Status::OK(); } +// used by adj_hessian_op. +tensorflow::Status GetProgramsCoefficients( + tensorflow::OpKernelContext* context, std::vector* programs_coeffs, + std::vector>* other_programs_coeffs) { + const Tensor* input_grads; + Status status = context->input("programs_coeffs", &input_grads); + if (!status.ok()) { + return status; + } + + if (input_grads->dims() != 1) { + return Status(tensorflow::error::INVALID_ARGUMENT, + absl::StrCat("programs_coeffs must be rank 1. Got rank ", + input_grads->dims(), ".")); + } + + const auto vec_coeff = input_grads->vec(); + programs_coeffs->reserve(vec_coeff.dimension(0)); + for (unsigned int i = 0; i < vec_coeff.dimension(0); i++) { + const float grad_v = vec_coeff(i); + programs_coeffs->push_back(grad_v); + } + + status = context->input("other_programs_coeffs", &input_grads); + if (!status.ok()) { + return status; + } + + if (input_grads->dims() != 2) { + return Status(tensorflow::error::INVALID_ARGUMENT, + absl::StrCat("other_programs_coeffs must be rank 2. " + "Got rank ", + input_grads->dims(), ".")); + } + + const auto matrix_coeffs = input_grads->matrix(); + other_programs_coeffs->reserve(matrix_coeffs.dimension(0)); + for (unsigned int i = 0; i < matrix_coeffs.dimension(0); i++) { + std::vector sub_coeffs; + sub_coeffs.reserve(matrix_coeffs.dimension(1)); + for (unsigned int j = 0; j < matrix_coeffs.dimension(1); j++) { + const float coeff = matrix_coeffs(i, j); + sub_coeffs.push_back(coeff); + } + other_programs_coeffs->push_back(sub_coeffs); + } + + return Status::OK(); +} + } // namespace tfq diff --git a/tensorflow_quantum/core/ops/parse_context.h b/tensorflow_quantum/core/ops/parse_context.h index e84f8e8bf..0ca1486b7 100644 --- a/tensorflow_quantum/core/ops/parse_context.h +++ b/tensorflow_quantum/core/ops/parse_context.h @@ -110,6 +110,11 @@ tensorflow::Status GetPrevGrads( tensorflow::OpKernelContext* context, std::vector>* parsed_prev_grads); +// Parses coefficients of programs & other_programs. Used by adjoint Hessian op. +tensorflow::Status GetProgramsCoefficients( + tensorflow::OpKernelContext* context, std::vector* programs_coeffs, + std::vector>* other_programs_coeffs); + } // namespace tfq #endif // TFQ_CORE_OPS_PARSE_CONTEXT diff --git a/tensorflow_quantum/core/src/BUILD b/tensorflow_quantum/core/src/BUILD index b0f69c95b..1771f3265 100644 --- a/tensorflow_quantum/core/src/BUILD +++ b/tensorflow_quantum/core/src/BUILD @@ -31,6 +31,23 @@ cc_library( ], ) +cc_library( + name = "adj_hessian_util", + srcs = ["adj_hessian_util.cc"], + hdrs = ["adj_hessian_util.h"], + deps = [ + ":adj_util", + ":circuit_parser_qsim", + "@qsim//lib:circuit", + "@qsim//lib:fuser", + "@qsim//lib:fuser_basic", + "@qsim//lib:gate", + "@qsim//lib:gates_cirq", + "@qsim//lib:io", + "@qsim//lib:matrix", + ], +) + cc_test( name = "adj_util_test", srcs = ["adj_util_test.cc"], diff --git a/tensorflow_quantum/core/src/adj_hessian_util.cc b/tensorflow_quantum/core/src/adj_hessian_util.cc new file mode 100644 index 000000000..0f5ce12e1 --- /dev/null +++ b/tensorflow_quantum/core/src/adj_hessian_util.cc @@ -0,0 +1,611 @@ +/* Copyright 2021 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 "tensorflow_quantum/core/src/adj_hessian_util.h" + +#include +#include +#include + +#include "../qsim/lib/circuit.h" +#include "../qsim/lib/fuser.h" +#include "../qsim/lib/fuser_basic.h" +#include "../qsim/lib/gate.h" +#include "../qsim/lib/gates_cirq.h" +#include "../qsim/lib/io.h" +#include "../qsim/lib/matrix.h" +#include "tensorflow_quantum/core/src/circuit_parser_qsim.h" + +namespace tfq { + +typedef qsim::Cirq::GateCirq QsimGate; +typedef qsim::Circuit QsimCircuit; + +void CreateHessianCircuit( + const QsimCircuit& circuit, const std::vector& metadata, + std::vector>>* partial_fuses, + std::vector* grad_gates) { + for (std::vector::size_type i = 0; i < metadata.size(); i++) { + if (metadata[i].symbol_values.empty()) { + continue; + } + // found a gate that was constructed with symbols. + GradientOfGate grad; + + // Single qubit Eigen. + if (circuit.gates[i].kind == qsim::Cirq::GateKind::kXPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kYPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kZPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kHPowGate) { + PopulateHessianSingleEigen( + metadata[i].create_f1, metadata[i].symbol_values[0], i, + circuit.gates[i].qubits[0], metadata[i].gate_params[0], + metadata[i].gate_params[1], metadata[i].gate_params[2], &grad); + grad_gates->push_back(grad); + } + + // Two qubit Eigen. + else if (circuit.gates[i].kind == qsim::Cirq::GateKind::kCZPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kCXPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kXXPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kYYPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kZZPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kISwapPowGate || + circuit.gates[i].kind == qsim::Cirq::GateKind::kSwapPowGate) { + bool swapq = circuit.gates[i].swapped; + PopulateHessianTwoEigen( + metadata[i].create_f2, metadata[i].symbol_values[0], i, + swapq ? circuit.gates[i].qubits[1] : circuit.gates[i].qubits[0], + swapq ? circuit.gates[i].qubits[0] : circuit.gates[i].qubits[1], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], &grad); + grad_gates->push_back(grad); + } + + // PhasedX + else if (circuit.gates[i].kind == qsim::Cirq::GateKind::kPhasedXPowGate) { + // Process potentially several symbols. + bool symbolic_pexp = false; + bool symbolic_exp = false; + for (std::vector >::size_type j = 0; j < metadata[i].symbol_values.size(); j++) { + if (metadata[i].placeholder_names[j] == + GateParamNames::kPhaseExponent) { + symbolic_pexp = true; + PopulateHessianPhasedXPhasedExponent( + metadata[i].symbol_values[j], i, circuit.gates[i].qubits[0], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], + metadata[i].gate_params[4], &grad); + } else if (metadata[i].placeholder_names[j] == + GateParamNames::kExponent) { + symbolic_exp = true; + PopulateHessianPhasedXExponent( + metadata[i].symbol_values[j], i, circuit.gates[i].qubits[0], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], + metadata[i].gate_params[4], &grad); + } + } + if (symbolic_pexp && symbolic_exp) { + PopulateCrossTermPhasedXPhasedExponentExponent( + i, circuit.gates[i].qubits[0], metadata[i].gate_params[0], + metadata[i].gate_params[1], metadata[i].gate_params[2], + metadata[i].gate_params[3], metadata[i].gate_params[4], &grad); + } + grad_gates->push_back(grad); + } + + // Fsim + else if (circuit.gates[i].kind == qsim::Cirq::GateKind::kFSimGate) { + // Process potentially several symbols. + + bool swapq = circuit.gates[i].swapped; + bool symbolic_theta = false; + bool symbolic_phi = false; + for (std::vector >::size_type j = 0; j < metadata[i].symbol_values.size(); j++) { + if (metadata[i].placeholder_names[j] == GateParamNames::kTheta) { + symbolic_theta = true; + PopulateHessianFsimTheta( + metadata[i].symbol_values[j], i, + swapq ? circuit.gates[i].qubits[1] : circuit.gates[i].qubits[0], + swapq ? circuit.gates[i].qubits[0] : circuit.gates[i].qubits[1], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], &grad); + } else if (metadata[i].placeholder_names[j] == GateParamNames::kPhi) { + symbolic_phi = true; + PopulateHessianFsimPhi( + metadata[i].symbol_values[j], i, + swapq ? circuit.gates[i].qubits[1] : circuit.gates[i].qubits[0], + swapq ? circuit.gates[i].qubits[0] : circuit.gates[i].qubits[1], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], &grad); + } + } + if (symbolic_theta && symbolic_phi) { + PopulateCrossTermFsimThetaPhi( + i, swapq ? circuit.gates[i].qubits[1] : circuit.gates[i].qubits[0], + swapq ? circuit.gates[i].qubits[0] : circuit.gates[i].qubits[1], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], &grad); + } + + grad_gates->push_back(grad); + } + + // PhasedISwap + else if (circuit.gates[i].kind == + qsim::Cirq::GateKind::kPhasedISwapPowGate) { + // Process potentially several symbols. + bool swapq = circuit.gates[i].swapped; + bool symbolic_pexp = false; + bool symbolic_exp = false; + for (std::vector >::size_type j = 0; j < metadata[i].symbol_values.size(); j++) { + if (metadata[i].placeholder_names[j] == + GateParamNames::kPhaseExponent) { + symbolic_pexp = true; + PopulateHessianPhasedISwapPhasedExponent( + metadata[i].symbol_values[j], i, + swapq ? circuit.gates[i].qubits[1] : circuit.gates[i].qubits[0], + swapq ? circuit.gates[i].qubits[0] : circuit.gates[i].qubits[1], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], &grad); + + } else if (metadata[i].placeholder_names[j] == + GateParamNames::kExponent) { + symbolic_exp = true; + PopulateHessianPhasedISwapExponent( + metadata[i].symbol_values[j], i, + swapq ? circuit.gates[i].qubits[1] : circuit.gates[i].qubits[0], + swapq ? circuit.gates[i].qubits[0] : circuit.gates[i].qubits[1], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], &grad); + } + } + if (symbolic_pexp && symbolic_exp) { + PopulateCrossTermPhasedISwapPhasedExponentExponent( + i, swapq ? circuit.gates[i].qubits[1] : circuit.gates[i].qubits[0], + swapq ? circuit.gates[i].qubits[0] : circuit.gates[i].qubits[1], + metadata[i].gate_params[0], metadata[i].gate_params[1], + metadata[i].gate_params[2], metadata[i].gate_params[3], &grad); + } + grad_gates->push_back(grad); + } + } + + // Produce partial fuses around the hessian gates. + auto fuser = qsim::BasicGateFuser(); + auto left = circuit.gates.begin(); + auto right = left; + + partial_fuses->assign(grad_gates->size() + 1, + std::vector>({})); + for (std::vector::size_type i = 0; i < grad_gates->size(); i++) { + right = circuit.gates.begin() + (*grad_gates)[i].index; + (*partial_fuses)[i] = + fuser.FuseGates(qsim::BasicGateFuser::Parameter(), + circuit.num_qubits, left, right); + left = right + 1; + } + right = circuit.gates.end(); + (*partial_fuses)[grad_gates->size()] = + fuser.FuseGates(qsim::BasicGateFuser::Parameter(), + circuit.num_qubits, left, right); +} + +void PopulateHessianSingleEigen( + const std::function& + create_f, + const std::string& symbol, unsigned int location, unsigned int qid, + float exp, float exp_s, float gs, GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = create_f(0, qid, (exp + _HESS_EPS) * exp_s, gs); + auto center = create_f(0, qid, exp * exp_s, gs); + auto right = create_f(0, qid, (exp - _HESS_EPS) * exp_s, gs); + // Due to precision issue, (1) multiplies weights first rather than last. + // and (2) doesn't use _INVERSE_HESS_EPS_SQUARE + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + std::cout << "left = ["; + int size = 8; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = left.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + std::cout << "right = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = right.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + std::cout << "center = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = center.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + Matrix2Add(right.matrix, + left.matrix); // left's entries have right added + std::cout << "left_plus_right = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = left.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto prefix = (val > 0) ? ((i % 2 == 1) ? "+" : "") : ""; + auto ending = (i % 2 == 0) ? "" : ((i == size - 1) ? "j" : "j,"); + std::cout << prefix << buf << ending; + } + std::cout << "]" << std::endl; + qsim::MatrixScalarMultiply(2.0, center.matrix); + std::cout << "twice_center = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = center.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + Matrix2Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + std::cout << "left_plus_right_minus_twice_center = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = left.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + grad->grad_gates.push_back(left); +} + +void PopulateHessianTwoEigen( + const std::function& create_f, + const std::string& symbol, unsigned int location, unsigned int qid, + unsigned int qid2, float exp, float exp_s, float gs, GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = create_f(0, qid, qid2, (exp + _HESS_EPS) * exp_s, gs); + auto center = create_f(0, qid, qid2, exp * exp_s, gs); + auto right = create_f(0, qid, qid2, (exp - _HESS_EPS) * exp_s, gs); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + std::cout << "PopulateHessianTwoEigen" << std::endl; + int size = 32; + std::cout << "left = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = left.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + std::cout << "center = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = center.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + std::cout << "right = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = right.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + Matrix4Add(right.matrix, + left.matrix); // left's entries have right added. + std::cout << "left_plus_right = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = left.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + qsim::MatrixScalarMultiply(2.0, center.matrix); + std::cout << "twice_center = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = center.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + Matrix4Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + std::cout << "left_plus_right_minus_twice_center = ["; + for (int i = 0; i < size; i++) { + char buf[100]; + auto val = left.matrix[i]; + std::sprintf(buf, "%.7e", val); + auto ending = ( + (i % 2 == 0) ? ((val > 0) ? "+" : "") : + ((i == size - 1) ? "j" : "j,")); + std::cout << buf << ending; + } + std::cout << "]" << std::endl; + grad->grad_gates.push_back(left); +} + +void PopulateHessianPhasedXPhasedExponent(const std::string& symbol, + unsigned int location, + unsigned int qid, float pexp, + float pexp_s, float exp, float exp_s, + float gs, GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, (pexp + _HESS_EPS) * pexp_s, exp * exp_s, gs); + auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, + exp * exp_s, gs); + auto right = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, (pexp - _HESS_EPS) * pexp_s, exp * exp_s, gs); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + Matrix2Add(right.matrix, + left.matrix); // left's entries have right added. + qsim::MatrixScalarMultiply(2.0, center.matrix); + Matrix2Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateHessianPhasedXExponent(const std::string& symbol, + unsigned int location, unsigned int qid, + float pexp, float pexp_s, float exp, + float exp_s, float gs, + GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, pexp * pexp_s, (exp + _HESS_EPS) * exp_s, gs); + auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, + exp * exp_s, gs); + auto right = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, pexp * pexp_s, (exp - _HESS_EPS) * exp_s, gs); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + Matrix2Add(right.matrix, + left.matrix); // left's entries have right added. + qsim::MatrixScalarMultiply(2.0, center.matrix); + Matrix2Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateCrossTermPhasedXPhasedExponentExponent( + unsigned int location, unsigned int qid, float pexp, float pexp_s, + float exp, float exp_s, float gs, GradientOfGate* grad) { + grad->params.push_back(kUsePrevTwoSymbols); + grad->index = location; + auto left = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); + auto left_center = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); + auto right_center = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); + auto right = qsim::Cirq::PhasedXPowGate::Create( + 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left_center.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right_center.matrix); + Matrix2Add(right.matrix, + left.matrix); // left's entries have right added. + Matrix2Add(right_center.matrix, left_center.matrix); + Matrix2Diff(left_center.matrix, + left.matrix); // left's entries have left_center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateHessianFsimTheta(const std::string& symbol, unsigned int location, + unsigned int qid, unsigned qid2, float theta, + float theta_s, float phi, float phi_s, + GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, (theta + _HESS_EPS) * theta_s, phi * phi_s); + auto center = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, theta * theta_s, phi * phi_s); + auto right = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, (theta - _HESS_EPS) * theta_s, phi * phi_s); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + Matrix4Add(right.matrix, + left.matrix); // left's entries have right added. + qsim::MatrixScalarMultiply(2.0, center.matrix); + Matrix4Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateHessianFsimPhi(const std::string& symbol, unsigned int location, + unsigned int qid, unsigned qid2, float theta, + float theta_s, float phi, float phi_s, + GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = qsim::Cirq::FSimGate::Create(0, qid, qid2, theta * theta_s, + (phi + _HESS_EPS) * phi_s); + auto center = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, theta * theta_s, phi * phi_s); + auto right = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, theta * theta_s, (phi - _HESS_EPS) * phi_s); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + Matrix4Add(right.matrix, + left.matrix); // left's entries have right added. + qsim::MatrixScalarMultiply(2.0, center.matrix); + Matrix4Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateCrossTermFsimThetaPhi(unsigned int location, unsigned int qid, + unsigned qid2, float theta, float theta_s, + float phi, float phi_s, + GradientOfGate* grad) { + grad->params.push_back(kUsePrevTwoSymbols); + grad->index = location; + auto left = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, (theta + _GRAD_EPS) * theta_s, (phi + _GRAD_EPS) * phi_s); + auto left_center = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, (theta + _GRAD_EPS) * theta_s, (phi - _GRAD_EPS) * phi_s); + auto right_center = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, (theta - _GRAD_EPS) * theta_s, (phi + _GRAD_EPS) * phi_s); + auto right = qsim::Cirq::FSimGate::Create( + 0, qid, qid2, (theta - _GRAD_EPS) * theta_s, (phi - _GRAD_EPS) * phi_s); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left_center.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right_center.matrix); + Matrix4Add(right.matrix, + left.matrix); // left's entries have right added. + Matrix4Add(right_center.matrix, left_center.matrix); + Matrix4Diff(left_center.matrix, + left.matrix); // left's entries have left_center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateHessianPhasedISwapPhasedExponent( + const std::string& symbol, unsigned int location, unsigned int qid, + unsigned int qid2, float pexp, float pexp_s, float exp, float exp_s, + GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, (pexp + _HESS_EPS) * pexp_s, exp * exp_s); + auto center = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, pexp * pexp_s, exp * exp_s); + auto right = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, (pexp - _HESS_EPS) * pexp_s, exp * exp_s); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + Matrix4Add(right.matrix, + left.matrix); // left's entries have right added. + qsim::MatrixScalarMultiply(2.0, center.matrix); + Matrix4Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateHessianPhasedISwapExponent(const std::string& symbol, + unsigned int location, unsigned int qid, + unsigned int qid2, float pexp, + float pexp_s, float exp, float exp_s, + GradientOfGate* grad) { + grad->params.push_back(symbol); + grad->index = location; + auto left = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, pexp * pexp_s, (exp + _HESS_EPS) * exp_s); + auto center = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, pexp * pexp_s, exp * exp_s); + auto right = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, pexp * pexp_s, (exp - _HESS_EPS) * exp_s); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); + Matrix4Add(right.matrix, + left.matrix); // left's entries have right added. + qsim::MatrixScalarMultiply(2.0, center.matrix); + Matrix4Diff(center.matrix, + left.matrix); // left's entries have center subtracted. + grad->grad_gates.push_back(left); +} + +void PopulateCrossTermPhasedISwapPhasedExponentExponent( + unsigned int location, unsigned int qid, unsigned int qid2, float pexp, + float pexp_s, float exp, float exp_s, GradientOfGate* grad) { + grad->params.push_back(kUsePrevTwoSymbols); + grad->index = location; + auto left = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, (pexp + _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s); + auto left_center = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, (pexp + _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s); + auto right_center = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, (pexp - _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s); + auto right = qsim::Cirq::PhasedISwapPowGate::Create( + 0, qid, qid2, (pexp - _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s); + // Due to precision issue, multiply weights first. + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left_center.matrix); + qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right_center.matrix); + Matrix4Add(right.matrix, + left.matrix); // left's entries have right added. + Matrix4Add(right_center.matrix, left_center.matrix); + Matrix4Diff(left_center.matrix, + left.matrix); // left's entries have center subtracted. + grad->grad_gates.push_back(left); +} + +} // namespace tfq diff --git a/tensorflow_quantum/core/src/adj_hessian_util.h b/tensorflow_quantum/core/src/adj_hessian_util.h new file mode 100644 index 000000000..0bc42c7c2 --- /dev/null +++ b/tensorflow_quantum/core/src/adj_hessian_util.h @@ -0,0 +1,124 @@ +/* Copyright 2021 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. +==============================================================================*/ + +#ifndef TFQ_CORE_SRC_ADJ_HESSIAN_UTIL_H_ +#define TFQ_CORE_SRC_ADJ_HESSIAN_UTIL_H_ + +#include +#include +#include + +#include "../qsim/lib/circuit.h" +#include "../qsim/lib/fuser.h" +#include "../qsim/lib/fuser_basic.h" +#include "../qsim/lib/gates_cirq.h" +#include "../qsim/lib/io.h" +#include "../qsim/lib/matrix.h" +#include "tensorflow_quantum/core/src/adj_util.h" +#include "tensorflow_quantum/core/src/circuit_parser_qsim.h" + +namespace tfq { + +static const float _HESS_EPS = 1e-2; +static const float _INVERSE_HESS_EPS_SQUARE = 1e4; +static const std::string kUsePrevTwoSymbols = "use_prev_two_symbols"; + +// Computes all gates who's hessian will need to be taken, in addition +// fuses all gates around those gates for faster circuit execution. +void CreateHessianCircuit( + const qsim::Circuit>& circuit, + const std::vector& metadata, + std::vector>>>* + partial_fuses, + std::vector* grad_gates); + +void PopulateHessianSingleEigen( + const std::function(unsigned int, unsigned int, + float, float)>& create_f, + const std::string& symbol, unsigned int location, unsigned int qid, + float exp, float exp_s, float gs, GradientOfGate* grad); + +void PopulateHessianTwoEigen( + const std::function( + unsigned int, unsigned int, unsigned int, float, float)>& create_f, + const std::string& symbol, unsigned int location, unsigned int qid, + unsigned int qid2, float exp, float exp_s, float gs, GradientOfGate* grad); + +// Note: all methods below expect gate qubit indices to have been swapped so +// qid < qid2. +void PopulateHessianPhasedXPhasedExponent(const std::string& symbol, + unsigned int location, + unsigned int qid, float pexp, + float pexp_s, float exp, float exp_s, + float gs, GradientOfGate* grad); + +void PopulateHessianPhasedXExponent(const std::string& symbol, + unsigned int location, unsigned int qid, + float pexp, float pexp_s, float exp, + float exp_s, float gs, + GradientOfGate* grad); + +void PopulateCrossTermPhasedXPhasedExponentExponent( + unsigned int location, unsigned int qid, float pexp, float pexp_s, + float exp, float exp_s, float gs, GradientOfGate* grad); + +void PopulateHessianFsimTheta(const std::string& symbol, unsigned int location, + unsigned int qid, unsigned qid2, float theta, + float theta_s, float phi, float phi_s, + GradientOfGate* grad); + +void PopulateHessianFsimPhi(const std::string& symbol, unsigned int location, + unsigned int qid, unsigned qid2, float theta, + float theta_s, float phi, float phi_s, + GradientOfGate* grad); + +void PopulateCrossTermFsimThetaPhi(unsigned int location, unsigned int qid, + unsigned qid2, float theta, float theta_s, + float phi, float phi_s, + GradientOfGate* grad); + +void PopulateHessianPhasedISwapPhasedExponent( + const std::string& symbol, unsigned int location, unsigned int qid, + unsigned int qid2, float pexp, float pexp_s, float exp, float exp_s, + GradientOfGate* grad); + +void PopulateHessianPhasedISwapExponent(const std::string& symbol, + unsigned int location, unsigned int qid, + unsigned int qid2, float pexp, + float pexp_s, float exp, float exp_s, + GradientOfGate* grad); + +void PopulateCrossTermPhasedISwapPhasedExponentExponent( + unsigned int location, unsigned int qid, unsigned int qid2, float pexp, + float pexp_s, float exp, float exp_s, GradientOfGate* grad); + +template +void Matrix2Add(Array2& source, Array2& dest) { + for (unsigned i = 0; i < 8; i++) { + dest[i] += source[i]; + } +} + +// does matrix elementwise addition dest += source. +template +void Matrix4Add(Array2& source, Array2& dest) { + for (unsigned i = 0; i < 32; i++) { + dest[i] += source[i]; + } +} + +} // namespace tfq + +#endif // TFQ_CORE_SRC_ADJ_UTIL_H_ diff --git a/tensorflow_quantum/core/src/adj_util.cc b/tensorflow_quantum/core/src/adj_util.cc index ceb76b2c1..c60bf7d57 100644 --- a/tensorflow_quantum/core/src/adj_util.cc +++ b/tensorflow_quantum/core/src/adj_util.cc @@ -29,8 +29,6 @@ limitations under the License. namespace tfq { -static const float _GRAD_EPS = 5e-3; - typedef qsim::Cirq::GateCirq QsimGate; typedef qsim::Circuit QsimCircuit; diff --git a/tensorflow_quantum/core/src/adj_util.h b/tensorflow_quantum/core/src/adj_util.h index 7cc383ad0..b9913f1de 100644 --- a/tensorflow_quantum/core/src/adj_util.h +++ b/tensorflow_quantum/core/src/adj_util.h @@ -30,6 +30,8 @@ limitations under the License. namespace tfq { +static const float _GRAD_EPS = 5e-3; + struct GradientOfGate { // name of parameters used by gate. std::vector params; @@ -100,7 +102,7 @@ void PopulateGradientPhasedISwapExponent(const std::string& symbol, float pexp, float pexp_s, float exp, float exp_s, GradientOfGate* grad); -// does matrix elementiwse subtraction dest -= source. +// does matrix elementwise subtraction dest -= source. template void Matrix2Diff(Array2& source, Array2& dest) { for (unsigned i = 0; i < 8; i++) { @@ -108,7 +110,7 @@ void Matrix2Diff(Array2& source, Array2& dest) { } } -// does matrix elementiwse subtraction dest -= source. +// does matrix elementwise subtraction dest -= source. template void Matrix4Diff(Array2& source, Array2& dest) { for (unsigned i = 0; i < 32; i++) { From 018e5411c141a6f55374919816fe3fbf3efa584f Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 06:52:09 +0900 Subject: [PATCH 02/12] Add analytic diff on PhasedXPowGate --- .../math_ops/inner_product_hessian_test.py | 89 +++++------- .../ops/math_ops/tfq_inner_product_hessian.cc | 49 +++---- .../core/src/adj_hessian_util.cc | 102 ++++++------- .../core/src/adj_hessian_util.h | 135 ++++++++++++++++++ 4 files changed, 244 insertions(+), 131 deletions(-) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py index 8d92b15df..c9459c872 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -47,7 +47,7 @@ cirq.FSimGate, ] -_ATOL_FOR_COMPLEX_GATE = 1e-1 +_ATOL_FOR_COMPLEX_GATE = 1e-2 _COMPLEX_GATES = [ cirq.PhasedXPowGate, ] @@ -72,11 +72,8 @@ def get_gate(gate, symbol_names, qubits): def get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, resolver): new_resolver = copy.deepcopy(resolver) - if name_j == name_k: - new_resolver.param_dict[name_j] += (dx_j + dx_k) - else: - new_resolver.param_dict[name_j] += dx_j - new_resolver.param_dict[name_k] += dx_k + new_resolver.param_dict[name_j] += dx_j + new_resolver.param_dict[name_k] += dx_k return cirq.resolve_parameters(circuit, new_resolver) @@ -91,8 +88,7 @@ def get_finite_difference_hessian(circuit, name_j, name_k, resolver): final_circuit_pm = get_shifted_resolved_circuit( circuit, name_j, name_k, dx, -dx, resolver) final_circuit_mm = get_shifted_resolved_circuit( - circuit, name_j, name_k, -dx, -dx, - resolver) + circuit, name_j, name_k, -dx, -dx, resolver) final_wf_pp = inv_square_two_dx * cirq.final_state_vector(final_circuit_pp) final_wf_mp = inv_square_two_dx * cirq.final_state_vector(final_circuit_mp) final_wf_pm = inv_square_two_dx * cirq.final_state_vector(final_circuit_pm) @@ -312,21 +308,21 @@ class InnerProductAdjHessianTest(tf.test.TestCase, parameterized.TestCase): 'batch_size': 1, 'inner_dim_size': 5 }, - # { - # 'n_qubits': 5, - # 'batch_size': 10, - # 'inner_dim_size': 1 - # }, - # { - # 'n_qubits': 10, - # 'batch_size': 10, - # 'inner_dim_size': 2 - # }, - # { - # 'n_qubits': 5, - # 'batch_size': 10, - # 'inner_dim_size': 5 - # }, + { + 'n_qubits': 5, + 'batch_size': 10, + 'inner_dim_size': 1 + }, + { + 'n_qubits': 10, + 'batch_size': 10, + 'inner_dim_size': 2 + }, + { + 'n_qubits': 5, + 'batch_size': 10, + 'inner_dim_size': 5 + }, ]) def correctness_with_symbols(self, n_qubits, batch_size, inner_dim_size): @@ -336,7 +332,7 @@ def correctness_with_symbols(self, n_qubits, batch_size, qubits = cirq.GridQubit.rect(1, n_qubits) circuit_batch, resolver_batch = \ util.random_symbol_circuit_resolver_batch( - qubits, symbol_names, batch_size, n_moments=2) + qubits, symbol_names, batch_size) print(circuit_batch) other_batch = [ @@ -494,33 +490,15 @@ class InnerProductHessianOnGates(tf.test.TestCase, parameterized.TestCase): @parameterized.parameters([ { - 'gate': [cirq.XPowGate], - 'symbol_names': ['alpha','beta','gamma']# names - }]) # for gate in _ONE_EIGEN_GATES + _TWO_EIGEN_GATES for names in _SYMBOL_NAMES]) + 'gate': gate, + 'symbol_names': names + } for gate in _ONE_EIGEN_GATES + _TWO_EIGEN_GATES + for names in _SYMBOL_NAMES]) def test_correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): """Tests that inner_product works with symbols.""" n_params = len(symbol_names) - qubits = cirq.GridQubit.rect(1, 5) # 2 if gate in _TWO_EIGEN_GATES else 1) - # circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] - circuit_batch = [cirq.Circuit([ - cirq.Moment( - (cirq.H**sympy.Mul(sympy.Float('0.96078327350981163', precision=53), sympy.Symbol('gamma'))).on(cirq.GridQubit(0, 0)), - (cirq.Y**sympy.Mul(sympy.Float('0.73193316007366105', precision=53), sympy.Symbol('alpha'))).on(cirq.GridQubit(0, 3)), - ), - cirq.Moment( - cirq.Y(cirq.GridQubit(0, 1)), - cirq.FSimGate(theta=0.123, phi=0.456).on(cirq.GridQubit(0, 4), cirq.GridQubit(0, 3)), - cirq.PhasedXPowGate(phase_exponent=0.123).on(cirq.GridQubit(0, 0)), - ), - cirq.Moment( - cirq.Y(cirq.GridQubit(0, 4)), - cirq.FSimGate(theta=0.123, phi=0.456).on(cirq.GridQubit(0, 2), cirq.GridQubit(0, 3)), - ), - cirq.Moment( - (cirq.H**sympy.Symbol('beta')).on(cirq.GridQubit(0, 0)), - ), - ])] - + qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) + circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] resolver_batch = [cirq.ParamResolver({name: 0.123 for name in symbol_names})] symbol_values_array = np.array( @@ -560,15 +538,12 @@ def test_correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): else: weighted_internal_wf += internal_wf for j, name_j in enumerate(symbol_names): - out_arr[i][j][j] = inner_product_op._inner_product_grad( - programs, symbol_names_tensor, symbol_values, other_programs, - other_programs_coeffs)[i][j] - # for k, name_k in enumerate(symbol_names): - # final_wf_grad = get_finite_difference_hessian( - # circuit_batch[i], name_j, name_k, resolver) - # out_arr[i][j][k] += ( - # programs_coeffs[i] * - # np.vdot(final_wf_grad, weighted_internal_wf)) + for k, name_k in enumerate(symbol_names): + final_wf_grad = get_finite_difference_hessian( + circuit_batch[i], name_j, name_k, resolver) + out_arr[i][j][k] += ( + programs_coeffs[i] * + np.vdot(final_wf_grad, weighted_internal_wf)) # Elapsed time should be less than 5% of cirq version. # (at least 20x speedup) diff --git a/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc index 427528258..6f176b884 100644 --- a/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc +++ b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc @@ -312,21 +312,6 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { std::cout << ">>>>>>... " << k << "th gradient gate is applied" << std::endl; qsim::ApplyGate(sim, hessian_gates[i][l - 1].grad_gates[k], scratch2); - auto ptr = scratch2.get(); - auto ptr_size = 2 << scratch2.num_qubits(); - std::cout << "Statevector" << std::endl; - for (int i = 0; i < ptr_size; i++) { - std::cout << ptr[i] << ","; - } - std::cout << std::endl; - - ptr = scratch.get(); - ptr_size = 2 << scratch.num_qubits(); - std::cout << "Other Statevector" << std::endl; - for (int i = 0; i < ptr_size; i++) { - std::cout << ptr[i] << ","; - } - std::cout << std::endl; // don't need not-found check since this is done upstream already. auto symbol = hessian_gates[i][l - 1].params[k]; std::cout << ">>>>>>... " << k << "th symbol = " << symbol << std::endl; @@ -389,7 +374,7 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { other_fused_circuits[i], sim, ss, scratch2, scratch); // now sv is |psi> - // scratch contains sum_j other_programs_coeffs[i][j]*|phi[i][j]> + // other_sv contains sum_j other_programs_coeffs[i][j]*|phi[i][j]> // Start adjoint differentiation on two gates // m is the index for the first gate std::cout << ">>> Start two gates hessian" << std::endl; @@ -414,30 +399,28 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { mask_m |= uint64_t{1} << control_loc; cbits_m |= ((cur_gate_m.cmask >> k) & 1) << control_loc; } + + ss.Copy(scratch, scratch4); + ss.Copy(sv, scratch2); for (std::vector::size_type p = 0; p < gradient_gates[i][m - 1].grad_gates.size(); p++) { // Copy sv onto scratch2 in anticipation of the first non-unitary // "gradient gate". - ss.Copy(sv, scratch2); if (!cur_gate_m.controlled_by.empty()) { // Gradient of controlled gates puts zeros on diagonal which is // the same as collapsing the state and then applying the // non-controlled version of the gradient gate. - ss.BulkSetAmpl(scratch2, mask_m, cbits_m, 0, 0, true); + ss.BulkSetAmpl(scratch4, mask_m, cbits_m, 0, 0, true); } std::cout << ">>>>>>(1)... p=" << p << "th gradient gate is applied" << std::endl; - qsim::ApplyGate(sim, gradient_gates[i][m - 1].grad_gates[p], - scratch2); + qsim::ApplyGateDagger(sim, gradient_gates[i][m - 1].grad_gates[p], + scratch4); // don't need not-found check since this is done upstream already. const auto it = maps[i].find(gradient_gates[i][m - 1].params[p]); std::cout << ">>>>>>(1)... p=" << p << "th symbol = " << gradient_gates[i][m - 1].params[p] << std::endl; const int loc_m = it->second.first; - // scratch2 is now (d/dsymbol[p])|psi> - // Copy scratch onto scratch4. - ss.Copy(scratch, scratch4); - // ApplyGateDagger(sim, cur_gate_m, scratch4); // n is the index for the second gate for (int n = m - 1; n >= 0; n--) { std::cout << ">>>>>>---(2) " << n << "th partial fused circuit is applied" << std::endl; @@ -455,6 +438,7 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { } // Hit a parameterized gate. + std::cout << "n-th gate index = " << gradient_gates[i][n - 1].index << std::endl; auto cur_gate_n = qsim_circuits[i].gates[gradient_gates[i][n - 1].index]; ApplyGateDagger(sim, cur_gate_n, scratch2); @@ -485,6 +469,21 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { qsim::ApplyGate(sim, gradient_gates[i][n - 1].grad_gates[q], scratch3); + auto ptr = scratch3.get(); + auto ptr_size = 2 << scratch3.num_qubits(); + std::cout << "Statevector" << std::endl; + for (int i = 0; i < ptr_size; i++) { + std::cout << ptr[i] << ","; + } + std::cout << std::endl; + + ptr = scratch4.get(); + ptr_size = 2 << scratch4.num_qubits(); + std::cout << "Other Statevector" << std::endl; + for (int i = 0; i < ptr_size; i++) { + std::cout << ptr[i] << ","; + } + std::cout << std::endl; // don't need not-found check since this is done upstream already. const auto it = maps[i].find(gradient_gates[i][n - 1].params[q]); std::cout << ">>>>>>---(2)... q=" << q << "th symbol = " << gradient_gates[i][n - 1].params[q] << std::endl; @@ -543,7 +542,6 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { Simulator sim = Simulator(tfq_for); StateSpace ss = StateSpace(tfq_for); auto sv = ss.Create(largest_nq); - auto sv_adj = ss.Create(largest_nq); auto scratch = ss.Create(largest_nq); auto scratch2 = ss.Create(largest_nq); auto scratch3 = ss.Create(largest_nq); @@ -560,7 +558,6 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { if (nq > largest_nq) { largest_nq = nq; sv = ss.Create(largest_nq); - sv_adj = ss.Create(largest_nq); scratch = ss.Create(largest_nq); scratch2 = ss.Create(largest_nq); } diff --git a/tensorflow_quantum/core/src/adj_hessian_util.cc b/tensorflow_quantum/core/src/adj_hessian_util.cc index 0f5ce12e1..eb3970fcf 100644 --- a/tensorflow_quantum/core/src/adj_hessian_util.cc +++ b/tensorflow_quantum/core/src/adj_hessian_util.cc @@ -389,21 +389,23 @@ void PopulateHessianPhasedXPhasedExponent(const std::string& symbol, float gs, GradientOfGate* grad) { grad->params.push_back(symbol); grad->index = location; - auto left = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, (pexp + _HESS_EPS) * pexp_s, exp * exp_s, gs); - auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, - exp * exp_s, gs); - auto right = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, (pexp - _HESS_EPS) * pexp_s, exp * exp_s, gs); - // Due to precision issue, multiply weights first. - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); - Matrix2Add(right.matrix, - left.matrix); // left's entries have right added. - qsim::MatrixScalarMultiply(2.0, center.matrix); - Matrix2Diff(center.matrix, - left.matrix); // left's entries have center subtracted. +// auto left = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, (pexp + _HESS_EPS) * pexp_s, exp * exp_s, gs); +// auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, +// exp * exp_s, gs); +// auto right = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, (pexp - _HESS_EPS) * pexp_s, exp * exp_s, gs); +// // Due to precision issue, multiply weights first. +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); +// Matrix2Add(right.matrix, +// left.matrix); // left's entries have right added. +// qsim::MatrixScalarMultiply(2.0, center.matrix); +// Matrix2Diff(center.matrix, +// left.matrix); // left's entries have center subtracted. + auto left = D2PhasedExponentPhasedXPowGate::Create( + 0, qid, pexp, pexp_s, exp*exp_s, gs); grad->grad_gates.push_back(left); } @@ -414,21 +416,23 @@ void PopulateHessianPhasedXExponent(const std::string& symbol, GradientOfGate* grad) { grad->params.push_back(symbol); grad->index = location; - auto left = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, pexp * pexp_s, (exp + _HESS_EPS) * exp_s, gs); - auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, - exp * exp_s, gs); - auto right = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, pexp * pexp_s, (exp - _HESS_EPS) * exp_s, gs); - // Due to precision issue, multiply weights first. - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); - Matrix2Add(right.matrix, - left.matrix); // left's entries have right added. - qsim::MatrixScalarMultiply(2.0, center.matrix); - Matrix2Diff(center.matrix, - left.matrix); // left's entries have center subtracted. +// auto left = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, pexp * pexp_s, (exp + _HESS_EPS) * exp_s, gs); +// auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, +// exp * exp_s, gs); +// auto right = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, pexp * pexp_s, (exp - _HESS_EPS) * exp_s, gs); +// // Due to precision issue, multiply weights first. +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); +// Matrix2Add(right.matrix, +// left.matrix); // left's entries have right added. +// qsim::MatrixScalarMultiply(2.0, center.matrix); +// Matrix2Diff(center.matrix, +// left.matrix); // left's entries have center subtracted. + auto left = D2ExponentPhasedXPowGate::Create( + 0, qid, pexp * pexp_s, exp, exp_s, gs); grad->grad_gates.push_back(left); } @@ -437,24 +441,26 @@ void PopulateCrossTermPhasedXPhasedExponentExponent( float exp, float exp_s, float gs, GradientOfGate* grad) { grad->params.push_back(kUsePrevTwoSymbols); grad->index = location; - auto left = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); - auto left_center = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); - auto right_center = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); - auto right = qsim::Cirq::PhasedXPowGate::Create( - 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); - // Due to precision issue, multiply weights first. - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left_center.matrix); - qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right_center.matrix); - Matrix2Add(right.matrix, - left.matrix); // left's entries have right added. - Matrix2Add(right_center.matrix, left_center.matrix); - Matrix2Diff(left_center.matrix, - left.matrix); // left's entries have left_center subtracted. +// auto left = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); +// auto left_center = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); +// auto right_center = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); +// auto right = qsim::Cirq::PhasedXPowGate::Create( +// 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); +// // Due to precision issue, multiply weights first. +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left_center.matrix); +// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right_center.matrix); +// Matrix2Add(right.matrix, +// left.matrix); // left's entries have right added. +// Matrix2Add(right_center.matrix, left_center.matrix); +// Matrix2Diff(left_center.matrix, +// left.matrix); // left's entries have left_center subtracted. + auto left = DPhasedExponentDExponentPhasedXPowGate::Create( + 0, qid, pexp, pexp_s, exp, exp_s, gs); grad->grad_gates.push_back(left); } diff --git a/tensorflow_quantum/core/src/adj_hessian_util.h b/tensorflow_quantum/core/src/adj_hessian_util.h index 0bc42c7c2..4ad665f95 100644 --- a/tensorflow_quantum/core/src/adj_hessian_util.h +++ b/tensorflow_quantum/core/src/adj_hessian_util.h @@ -119,6 +119,141 @@ void Matrix4Add(Array2& source, Array2& dest) { } } +// Due to the large error from the finite differencing PhasedXPowGate, +// Here we introduce analytically differentiated version of PhasedXPowGate up to +// 2nd order with double precision. + +template +struct D2PhasedExponentPhasedXPowGate { + static constexpr qsim::Cirq::GateKind kind = qsim::Cirq::kPhasedXPowGate; + static constexpr char name[] = "D2PhasedExponentPhasedXPowGate"; + static constexpr unsigned num_qubits = 1; + static constexpr bool symmetric = true; + + static qsim::Cirq::GateCirq Create(unsigned time, unsigned q0, + fp_type phase_exponent, + fp_type phase_exponent_scalar, + fp_type exponent = 1, + fp_type global_shift = 0) { + double pi = qsim::Cirq::pi_double; + double pexp = static_cast(phase_exponent); + double pexp_s = static_cast(phase_exponent_scalar); + double exp = static_cast(exponent); + double ec = std::cos(pi * exp); + double es = std::sin(pi * exp); + double gc = std::cos(pi * exp * global_shift); + double gs = std::sin(pi * exp * global_shift); + + double d2p_pc = -pi*pi*pexp_s*pexp_s*std::cos(pi * pexp_s * pexp); + double d2p_ps = -pi*pi*pexp_s*pexp_s*std::sin(pi * pexp_s * pexp); + + double br = -0.5 * ((-1 + ec) * gc - es * gs); + double bi = -0.5 * ((-1 + ec) * gs + es * gc); + + return qsim::CreateGate, D2PhasedExponentPhasedXPowGate>( + time, {q0}, {0., 0., static_cast(d2p_pc * br + d2p_ps * bi), + static_cast(d2p_pc * bi - d2p_ps * br), + static_cast(d2p_pc * br - d2p_ps * bi), + static_cast(d2p_pc * bi + d2p_ps * br), 0., 0.}, + {phase_exponent, exponent, global_shift}); + } +}; + +template +struct D2ExponentPhasedXPowGate { + static constexpr qsim::Cirq::GateKind kind = qsim::Cirq::kPhasedXPowGate; + static constexpr char name[] = "D2ExponentPhasedXPowGate"; + static constexpr unsigned num_qubits = 1; + static constexpr bool symmetric = true; + + static constexpr fp_type pi = static_cast(qsim::Cirq::pi_double); + + static qsim::Cirq::GateCirq Create(unsigned time, unsigned q0, + fp_type phase_exponent, + fp_type exponent = 1, + fp_type exponent_scalar = 1, + fp_type global_shift = 0) { + double pi = qsim::Cirq::pi_double; + double pexp = static_cast(phase_exponent); + double exp = static_cast(exponent); + double exp_s = static_cast(exponent_scalar); + double pc = std::cos(pi * pexp); + double ps = std::sin(pi * pexp); + double ec = std::cos(pi * exp_s * exp); + double es = std::sin(pi * exp_s * exp); + double gc = std::cos(pi * exp_s * exp * global_shift); + double gs = std::sin(pi * exp_s * exp * global_shift); + double dec = -pi * exp_s * std::sin(pi * exp_s * exp); + double des = pi * exp_s * std::cos(pi * exp_s * exp); + double dgc = -pi * exp_s * global_shift * std::sin(pi * exp_s * exp * global_shift); + double dgs = pi * exp_s * global_shift * std::cos(pi * exp_s * exp * global_shift); + double d2ec = - pi * exp_s * pi * exp_s * std::cos(pi * exp_s * exp); + double d2es = - pi * exp_s * pi * exp_s * std::sin(pi * exp_s * exp); + double d2gc = - pi * exp_s * pi * exp_s * global_shift * global_shift * std::cos(pi * exp_s * exp * global_shift); + double d2gs = - pi * exp_s * pi * exp_s * global_shift * global_shift * std::sin(pi * exp_s * exp * global_shift); + double common_r_front = ec * d2gc + 2.0 * dec * dgc + d2ec * gc; + double common_r_back = d2es * gs + 2.0 * des * dgs + es * d2gs; + double common_i_front = ec * d2gs + 2.0 * dec * dgs + d2ec * gs; + double common_i_back = d2es * gc + 2.0 * des * dgc + es * d2gc; + + fp_type d2ar = static_cast(0.5 * (d2gc + common_r_front - common_r_back)); + fp_type d2ai = static_cast(0.5 * (d2gs + common_i_front + common_i_back)); + double d2br = -0.5 * (-d2gc + common_r_front - common_r_back); + double d2bi = -0.5 * (-d2gs + common_i_front + common_i_back); + + return qsim::CreateGate, D2ExponentPhasedXPowGate>( + time, {q0}, {d2ar, d2ai, static_cast(pc * d2br + ps * d2bi), + static_cast(pc * d2bi - ps * d2br), + static_cast(pc * d2br - ps * d2bi), + static_cast(pc * d2bi + ps * d2br), d2ar, d2ai}, + {phase_exponent, exponent, global_shift}); + } +}; + +template +struct DPhasedExponentDExponentPhasedXPowGate { + static constexpr qsim::Cirq::GateKind kind = qsim::Cirq::kPhasedXPowGate; + static constexpr char name[] = "DPhasedExponentDExponentPhasedXPowGate"; + static constexpr unsigned num_qubits = 1; + static constexpr bool symmetric = true; + + static constexpr fp_type pi = static_cast(qsim::Cirq::pi_double); + + static qsim::Cirq::GateCirq Create(unsigned time, unsigned q0, + fp_type phase_exponent, + fp_type phase_exponent_scalar, + fp_type exponent = 1, + fp_type exponent_scalar = 1, + fp_type global_shift = 0) { + double pi = qsim::Cirq::pi_double; + double pexp = static_cast(phase_exponent); + double pexp_s = static_cast(phase_exponent_scalar); + double exp = static_cast(exponent); + double exp_s = static_cast(exponent_scalar); + double ec = std::cos(pi * exp_s * exp); + double es = std::sin(pi * exp_s * exp); + double gc = std::cos(pi * exp_s * exp * global_shift); + double gs = std::sin(pi * exp_s * exp * global_shift); + + double dp_pc = -pi * pexp_s * std::sin(pi * pexp_s * pexp); + double dp_ps = pi * pexp_s * std::cos(pi * pexp_s * pexp); + double de_ec = -pi * exp_s * std::sin(pi * exp_s * exp); + double de_es = pi * exp_s * std::cos(pi * exp_s * exp); + double de_gc = -pi * exp_s * global_shift * std::sin(pi * exp_s * global_shift * exp); + double de_gs = pi * exp_s * global_shift * std::cos(pi * exp_s * global_shift * exp); + + double de_br = -0.5 * ((-1 + ec) * de_gc + de_ec * gc - de_es * gs - es * de_gs); + double de_bi = -0.5 * ((-1 + ec) * de_gs + de_ec * gs + de_es * gc + es * de_gc); + + return qsim::CreateGate, DPhasedExponentDExponentPhasedXPowGate>( + time, {q0}, {0.,0., static_cast(dp_pc * de_br + dp_ps * de_bi), + static_cast(dp_pc * de_bi - dp_ps * de_br), + static_cast(dp_pc * de_br - dp_ps * de_bi), + static_cast(dp_pc * de_bi + dp_ps * de_br), 0.,0.}, + {phase_exponent, exponent, global_shift}); + } +}; + } // namespace tfq #endif // TFQ_CORE_SRC_ADJ_UTIL_H_ From 0079d53f963ad1ad1d1c11441b0bbce9e5ddd4a8 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 09:17:49 +0900 Subject: [PATCH 03/12] Fix the adjoint hessian bug and add ComputeSmall --- .../math_ops/inner_product_hessian_test.py | 323 ++++++++++-------- .../core/ops/math_ops/inner_product_op.py | 6 +- .../ops/math_ops/tfq_inner_product_hessian.cc | 311 ++++++++++------- .../core/src/adj_hessian_util.cc | 261 +------------- .../core/src/adj_hessian_util.h | 146 -------- tensorflow_quantum/python/util.py | 19 +- 6 files changed, 390 insertions(+), 676 deletions(-) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py index c9459c872..707ad1fa9 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -26,14 +26,14 @@ from tensorflow_quantum.python import util _INVERSE_SPEEDUP = 1 / 20.0 -_ATOL = 3e-3 +_ATOL = 0.2 +_RTOL = 0.3 _SYMBOL_NAMES = [['alpha'], ['alpha', 'beta']] _ONE_EIGEN_GATES = [ cirq.XPowGate, cirq.YPowGate, cirq.ZPowGate, cirq.HPowGate, - cirq.PhasedXPowGate, ] _TWO_EIGEN_GATES = [ cirq.XXPowGate, @@ -46,9 +46,7 @@ cirq.PhasedISwapPowGate, cirq.FSimGate, ] - -_ATOL_FOR_COMPLEX_GATE = 1e-2 -_COMPLEX_GATES = [ +_UNSUPPORTED_GATES = [ cirq.PhasedXPowGate, ] @@ -60,14 +58,20 @@ def get_gate(gate, symbol_names, qubits): else: a, b = symbols if gate == cirq.PhasedXPowGate or gate == cirq.PhasedISwapPowGate: - return [gate(phase_exponent=0.1 * a, exponent=0.2*b).on(*qubits), - gate(phase_exponent=0.3 * b, exponent=0.4*a).on(*qubits)] + return [ + gate(phase_exponent=0.1 * a, exponent=0.2 * b).on(*qubits), + gate(phase_exponent=0.3 * b, exponent=0.4 * a).on(*qubits) + ] if gate == cirq.FSimGate: - return [gate(theta=0.1 * a, phi=0.2*b).on(*qubits), - gate(theta=0.3 * b, phi=0.4*a).on(*qubits)] + return [ + gate(theta=0.1 * a, phi=0.2 * b).on(*qubits), + gate(theta=0.3 * b, phi=0.4 * a).on(*qubits) + ] - return [gate(exponent=0.1 * a).on(*qubits), - gate(exponent=0.2 * b).on(*qubits)] + return [ + gate(exponent=0.1 * a).on(*qubits), + gate(exponent=0.2 * b).on(*qubits) + ] def get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, resolver): @@ -80,15 +84,15 @@ def get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, resolver): def get_finite_difference_hessian(circuit, name_j, name_k, resolver): # dx came from _GRAD_EPS of core/src/adj_util.cc dx = 5e-3 - inv_square_two_dx = np.asarray([1e4+0.j], dtype=np.complex64) - final_circuit_pp = get_shifted_resolved_circuit( - circuit, name_j, name_k, dx, dx, resolver) - final_circuit_mp = get_shifted_resolved_circuit( - circuit, name_j, name_k, -dx, dx, resolver) - final_circuit_pm = get_shifted_resolved_circuit( - circuit, name_j, name_k, dx, -dx, resolver) - final_circuit_mm = get_shifted_resolved_circuit( - circuit, name_j, name_k, -dx, -dx, resolver) + inv_square_two_dx = np.asarray([1e4 + 0.j], dtype=np.complex64) + final_circuit_pp = get_shifted_resolved_circuit(circuit, name_j, name_k, dx, + dx, resolver) + final_circuit_mp = get_shifted_resolved_circuit(circuit, name_j, name_k, + -dx, dx, resolver) + final_circuit_pm = get_shifted_resolved_circuit(circuit, name_j, name_k, dx, + -dx, resolver) + final_circuit_mm = get_shifted_resolved_circuit(circuit, name_j, name_k, + -dx, -dx, resolver) final_wf_pp = inv_square_two_dx * cirq.final_state_vector(final_circuit_pp) final_wf_mp = inv_square_two_dx * cirq.final_state_vector(final_circuit_mp) final_wf_pm = inv_square_two_dx * cirq.final_state_vector(final_circuit_pm) @@ -324,16 +328,15 @@ class InnerProductAdjHessianTest(tf.test.TestCase, parameterized.TestCase): 'inner_dim_size': 5 }, ]) - def correctness_with_symbols(self, n_qubits, batch_size, + def test_correctness_with_symbols(self, n_qubits, batch_size, inner_dim_size): """Tests that inner_product works with symbols.""" - symbol_names = ['alpha', 'beta', 'gamma'] + symbol_names = ['alpha', 'beta', 'gamma', 'delta', 'eta', 'kappa'] n_params = len(symbol_names) qubits = cirq.GridQubit.rect(1, n_qubits) circuit_batch, resolver_batch = \ util.random_symbol_circuit_resolver_batch( - qubits, symbol_names, batch_size) - print(circuit_batch) + qubits, symbol_names, batch_size, exclude_gates=_UNSUPPORTED_GATES) other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] @@ -381,152 +384,146 @@ def correctness_with_symbols(self, n_qubits, batch_size, # Elapsed time should be less than 5% of cirq version. # (at least 20x speedup) - self.assertLess(t1, t2*_INVERSE_SPEEDUP) - self.assertAllClose(out, out_arr, atol=_ATOL) + self.assertLess(t1, t2 * _INVERSE_SPEEDUP) + self.assertAllClose(out, out_arr, atol=_ATOL, rtol=_RTOL) - # - # @parameterized.parameters([ - # { - # 'n_qubits': 5, - # 'batch_size': 1, - # 'inner_dim_size': 5 - # }, - # { - # 'n_qubits': 5, - # 'batch_size': 10, - # 'inner_dim_size': 1 - # }, - # { - # 'n_qubits': 10, - # 'batch_size': 10, - # 'inner_dim_size': 2 - # }, - # { - # 'n_qubits': 5, - # 'batch_size': 10, - # 'inner_dim_size': 5 - # }, - # ]) - # def test_correctness_without_symbols(self, n_qubits, batch_size, - # inner_dim_size): - # """Tests that inner_product_adj_grad works without symbols.""" - # qubits = cirq.GridQubit.rect(1, n_qubits) - # circuit_batch, _ = \ - # util.random_circuit_resolver_batch( - # qubits, batch_size) - # - # other_batch = [ - # util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] - # for _ in range(batch_size) - # ] - # - # programs = util.convert_to_tensor(circuit_batch) - # other_programs = util.convert_to_tensor(other_batch) - # symbol_names = tf.convert_to_tensor([], dtype=tf.dtypes.string) - # symbol_values = tf.convert_to_tensor([[] for _ in range(batch_size)]) - # progams_coeffs = np.ones((batch_size,)) - # other_programs_coeffs = np.ones((batch_size, inner_dim_size)) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'symbols must be a positive integer'): - # inner_product_op.inner_product_hessian(programs, symbol_names, - # symbol_values, - # other_programs, - # progams_coeffs, - # other_programs_coeffs) - - # def test_correctness_empty(self): - # """Tests the inner product adj grad between two empty circuits.""" - # symbol_names = ['alpha', 'beta'] - # empty_cicuit = util.convert_to_tensor([cirq.Circuit()]) - # empty_symbols = tf.convert_to_tensor([], dtype=tf.dtypes.string) - # empty_values = tf.convert_to_tensor([[]]) - # other_program = util.convert_to_tensor([[cirq.Circuit()]]) - # program_coeffs = np.ones((1,)) - # other_program_coeffs = np.ones((1, 1)) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'symbols must be a positive integer'): - # inner_product_op.inner_product_hessian(empty_cicuit, empty_symbols, - # empty_values, other_program, - # program_coeffs, - # other_program_coeffs) - # - # empty_cicuit = util.convert_to_tensor([cirq.Circuit()]) - # symbol_names = tf.convert_to_tensor(symbol_names, - # dtype=tf.dtypes.string) - # symbol_values = tf.convert_to_tensor([[0.0 for _ in range(2)]]) - # other_program = util.convert_to_tensor([[cirq.Circuit()]]) - # - # out = inner_product_op.inner_product_hessian(empty_cicuit, symbol_names, - # symbol_values, - # other_program, - # program_coeffs, - # other_program_coeffs) - # expected = np.zeros((1, len(symbol_names)), dtype=np.complex64) - # self.assertAllClose(out, expected) - - # def test_correctness_no_circuit(self): - # """Test the inner product grad between no 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) - # other_program = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.string) - # empty_program_coeffs = tf.raw_ops.Empty(shape=(0,), dtype=tf.float32) - # empty_other_program_coeffs = tf.raw_ops.Empty(shape=(0, 0), - # dtype=tf.float32) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'number of symbols must be a positive'): - # # When using `tf.gradients`, a user will never encounter this error - # # thanks to the `tf.cond` inside of the custom gradient. - # _ = inner_product_op.inner_product_hessian( - # empty_circuit, empty_symbols, empty_values, other_program, - # empty_program_coeffs, empty_other_program_coeffs) + @parameterized.parameters([ + { + 'n_qubits': 5, + 'batch_size': 1, + 'inner_dim_size': 5 + }, + { + 'n_qubits': 5, + 'batch_size': 10, + 'inner_dim_size': 1 + }, + { + 'n_qubits': 10, + 'batch_size': 10, + 'inner_dim_size': 2 + }, + { + 'n_qubits': 5, + 'batch_size': 10, + 'inner_dim_size': 5 + }, + ]) + def correctness_without_symbols(self, n_qubits, batch_size, inner_dim_size): + """Tests that inner_product_adj_grad works without symbols.""" + qubits = cirq.GridQubit.rect(1, n_qubits) + circuit_batch, _ = \ + util.random_circuit_resolver_batch( + qubits, batch_size) + + other_batch = [ + util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] + for _ in range(batch_size) + ] + + programs = util.convert_to_tensor(circuit_batch) + other_programs = util.convert_to_tensor(other_batch) + symbol_names = tf.convert_to_tensor([], dtype=tf.dtypes.string) + symbol_values = tf.convert_to_tensor([[] for _ in range(batch_size)]) + progams_coeffs = np.ones((batch_size,)) + other_programs_coeffs = np.ones((batch_size, inner_dim_size)) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbols must be a positive integer'): + inner_product_op.inner_product_hessian(programs, symbol_names, + symbol_values, + other_programs, + progams_coeffs, + other_programs_coeffs) + + def correctness_empty(self): + """Tests the inner product adj grad between two empty circuits.""" + symbol_names = ['alpha', 'beta'] + n_params = len(symbol_names) + empty_cicuit = util.convert_to_tensor([cirq.Circuit()]) + empty_symbols = tf.convert_to_tensor([], dtype=tf.dtypes.string) + empty_values = tf.convert_to_tensor([[]]) + other_program = util.convert_to_tensor([[cirq.Circuit()]]) + program_coeffs = np.ones((1,)) + other_program_coeffs = np.ones((1, 1)) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbols must be a positive integer'): + inner_product_op.inner_product_hessian(empty_cicuit, empty_symbols, + empty_values, other_program, + program_coeffs, + other_program_coeffs) + + empty_cicuit = util.convert_to_tensor([cirq.Circuit()]) + symbol_names = tf.convert_to_tensor(symbol_names, + dtype=tf.dtypes.string) + symbol_values = tf.convert_to_tensor([[0.0 for _ in range(2)]]) + other_program = util.convert_to_tensor([[cirq.Circuit()]]) + + out = inner_product_op.inner_product_hessian(empty_cicuit, symbol_names, + symbol_values, + other_program, + program_coeffs, + other_program_coeffs) + expected = np.zeros((1, n_params, n_params), dtype=np.complex64) + self.assertAllClose(out, expected) + + def correctness_no_circuit(self): + """Test the inner product grad between no 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) + other_program = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.string) + empty_program_coeffs = tf.raw_ops.Empty(shape=(0,), dtype=tf.float32) + empty_other_program_coeffs = tf.raw_ops.Empty(shape=(0, 0), + dtype=tf.float32) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'number of symbols must be a positive'): + # When using `tf.gradients`, a user will never encounter this error + # thanks to the `tf.cond` inside of the custom gradient. + _ = inner_product_op.inner_product_hessian( + empty_circuit, empty_symbols, empty_values, other_program, + empty_program_coeffs, empty_other_program_coeffs) class InnerProductHessianOnGates(tf.test.TestCase, parameterized.TestCase): - @parameterized.parameters([ - { - 'gate': gate, - 'symbol_names': names - } for gate in _ONE_EIGEN_GATES + _TWO_EIGEN_GATES - for names in _SYMBOL_NAMES]) - def test_correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): + @parameterized.parameters([{ + 'gate': gate, + 'symbol_names': names + } + for gate in _ONE_EIGEN_GATES + _TWO_EIGEN_GATES + for names in _SYMBOL_NAMES]) + def correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): """Tests that inner_product works with symbols.""" n_params = len(symbol_names) qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] - resolver_batch = [cirq.ParamResolver({name: 0.123 for name in symbol_names})] + resolver_batch = [ + cirq.ParamResolver({name: 0.123 for name in symbol_names}) + ] symbol_values_array = np.array( [[resolver[symbol] for symbol in symbol_names] for resolver in resolver_batch]) other_batch = [ - [cirq.Circuit(cirq.H.on_each(*qubits))] - for _ in circuit_batch + [cirq.Circuit(cirq.H.on_each(*qubits))] for _ in circuit_batch ] programs = util.convert_to_tensor(circuit_batch) other_programs = util.convert_to_tensor(other_batch) symbol_names_tensor = tf.convert_to_tensor(symbol_names, dtype=tf.dtypes.string) symbol_values = tf.convert_to_tensor(symbol_values_array) - programs_coeffs = tf.cast(tf.ones((1,)), tf.complex64) - other_programs_coeffs = tf.cast(tf.ones((1, 1)), tf.complex64) - # programs_coeffs = tf.cast(tf.random.normal((1,)), tf.complex64) - # other_programs_coeffs = tf.cast( - # tf.random.normal((1, 1)), tf.complex64) + programs_coeffs = tf.cast(tf.random.normal((1,)), tf.complex64) + other_programs_coeffs = tf.cast(tf.random.normal((1, 1)), tf.complex64) - t1 = time.time() out = inner_product_op.inner_product_hessian( programs, symbol_names_tensor, symbol_values, other_programs, programs_coeffs, other_programs_coeffs) - t1 = time.time() - t1 - t2 = time.time() out_arr = np.zeros((1, n_params, n_params), dtype=np.complex64) for i, resolver in enumerate(resolver_batch): weighted_internal_wf = None @@ -545,10 +542,40 @@ def test_correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): programs_coeffs[i] * np.vdot(final_wf_grad, weighted_internal_wf)) - # Elapsed time should be less than 5% of cirq version. - # (at least 20x speedup) - self.assertLess(t1, t2*_INVERSE_SPEEDUP) - self.assertAllClose(out, out_arr, atol=_ATOL_FOR_COMPLEX_GATE if gate in _COMPLEX_GATES else _ATOL) + self.assertAllClose(out, out_arr, atol=_ATOL, rtol=_RTOL) + + @parameterized.parameters([{ + 'gate': gate, + } for gate in _UNSUPPORTED_GATES for names in _SYMBOL_NAMES]) + def unsupported_gate_with_symbols(self, gate): + """Tests that inner_product works with symbols.""" + symbol_names = ['alpha'] + qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) + circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] + resolver_batch = [ + cirq.ParamResolver({name: 0.123 for name in symbol_names}) + ] + + symbol_values_array = np.array( + [[resolver[symbol] + for symbol in symbol_names] + for resolver in resolver_batch]) + other_batch = [ + [cirq.Circuit(cirq.H.on_each(*qubits))] for _ in circuit_batch + ] + programs = util.convert_to_tensor(circuit_batch) + other_programs = util.convert_to_tensor(other_batch) + symbol_names_tensor = tf.convert_to_tensor(symbol_names, + dtype=tf.dtypes.string) + symbol_values = tf.convert_to_tensor(symbol_values_array) + programs_coeffs = tf.cast(tf.random.normal((1,)), tf.complex64) + other_programs_coeffs = tf.cast(tf.random.normal((1, 1)), tf.complex64) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'is currently not supported'): + _ = inner_product_op.inner_product_hessian( + programs, symbol_names_tensor, symbol_values, other_programs, + programs_coeffs, other_programs_coeffs) if __name__ == "__main__": diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py index 0e2c8d961..db82bd999 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py @@ -67,9 +67,9 @@ def inner_product_hessian(programs, symbol_names, symbol_values, other_programs, (other_)programs[i] is weighted by (other_)programs_coeffs[i]. """ return MATH_OP_MODULE.tfq_inner_product_hessian( - programs, symbol_names, tf.cast(symbol_values, tf.float32), - other_programs, tf.cast(programs_coeffs, tf.float32), - tf.cast(other_programs_coeffs, tf.float32)) + programs, symbol_names, tf.cast(symbol_values, tf.float32), + other_programs, tf.cast(programs_coeffs, tf.float32), + tf.cast(other_programs_coeffs, tf.float32)) def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, diff --git a/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc index 6f176b884..0930d4d25 100644 --- a/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc +++ b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc @@ -15,7 +15,6 @@ limitations under the License. #include #include -#include #include "../qsim/lib/circuit.h" #include "../qsim/lib/gate_appl.h" @@ -134,6 +133,20 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { &fused_circuits[i], &gate_meta[i]); NESTED_FN_STATUS_SYNC(parse_status, local, p_lock); + for (std::vector::size_type j = 0; + j < gate_meta[i].size(); j++) { + if (gate_meta[i][j].symbol_values.empty()) { + continue; + } + if (qsim_circuits[i].gates[j].kind == + qsim::Cirq::GateKind::kPhasedXPowGate) { + NESTED_FN_STATUS_SYNC(parse_status, + Status(tensorflow::error::INVALID_ARGUMENT, + "the circuit with PhasedXPowGate is " + "currently not supported."), + p_lock); + } + } CreateGradientCircuit(qsim_circuits[i], gate_meta[i], &partial_fused_grad_circuits[i], &gradient_gates[i]); @@ -193,10 +206,10 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // Cross reference with standard google cloud compute instances // Memory ~= 2 * num_threads * (2 * 64 * 2 ** num_qubits in circuits) - // e2s2 = 2 CPU, 8GB -> Can safely do 22 since Memory = 4GB - // e2s4 = 4 CPU, 16GB -> Can safely do 22 since Memory = 8GB + // e2s2 = 2 CPU, 8GB -> Can safely do 20 since Memory = 4GB + // e2s4 = 4 CPU, 16GB -> Can safely do 20 since Memory = 8GB // ... - if (true || max_num_qubits >= 22 || output_dim_batch_size == 1) { + if (max_num_qubits >= 20 || output_dim_batch_size == 1) { ComputeLarge(num_qubits, maps, qsim_circuits, fused_circuits, partial_fused_grad_circuits, partial_fused_hess_circuits, gradient_gates, hessian_gates, other_fused_circuits, @@ -270,9 +283,7 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // now sv is |psi> // scratch contains sum_j other_programs_coeffs[i][j]*|phi[i][j]> // Start adjoint differentiation on a single gate - std::cout << ">>> Start a single gate hessian" << std::endl; for (int l = partial_fused_hess_circuits[i].size() - 1; l >= 0; l--) { - std::cout << ">>>>>> " << l << "th partial fused circuit is applied" << std::endl; for (int k = partial_fused_hess_circuits[i][l].size() - 1; k >= 0; k--) { ApplyFusedGateDagger(sim, partial_fused_hess_circuits[i][l][k], sv); @@ -309,19 +320,16 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // non-controlled version of the gradient gate. ss.BulkSetAmpl(scratch2, mask, cbits, 0, 0, true); } - std::cout << ">>>>>>... " << k << "th gradient gate is applied" << std::endl; qsim::ApplyGate(sim, hessian_gates[i][l - 1].grad_gates[k], scratch2); // don't need not-found check since this is done upstream already. auto symbol = hessian_gates[i][l - 1].params[k]; - std::cout << ">>>>>>... " << k << "th symbol = " << symbol << std::endl; double coeff = static_cast(programs_coeffs[i]); if (symbol == kUsePrevTwoSymbols) { // Apply second-order finite difference w.r.t. two symbols // That is, CrossTerm w.r.t. two symbols in one gate. auto symbol1 = hessian_gates[i][l - 1].params[k - 2]; auto symbol2 = hessian_gates[i][l - 1].params[k - 1]; - std::cout << ">>>>>>... CrossTerm : two symbols are " << symbol1 << " and " << symbol2 << std::endl; auto it = maps[i].find(symbol1); const int loc1 = it->second.first; it = maps[i].find(symbol2); @@ -332,17 +340,13 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // parameter-shift we need to apply a single `gradient_gate` // per a symbol. std::complex result = ss.InnerProduct(scratch2, scratch); - auto val = ( - std::complex(static_cast(coeff * result.real()), - static_cast(coeff * result.imag()))); + auto val = (std::complex( + static_cast(coeff * result.real()), + static_cast(coeff * result.imag()))); // Because Hessian is symmetric. - std::cout << ">>>>>>... value" << val << std::endl; (*output_tensor)(i, loc1, loc2) += val; - std::cout << ">>>>>>... (*output_tensor)(i, loc1, loc2)" << (*output_tensor)(i, loc1, loc2) << std::endl; (*output_tensor)(i, loc2, loc1) += val; - std::cout << ">>>>>>... (*output_tensor)(i, loc2, loc1)" << (*output_tensor)(i, loc2, loc1) << std::endl; } else { - std::cout << ">>>>>>... One gate " << symbol << std::endl; // Apply second-order finite difference w.r.t. one symbol const auto it = maps[i].find(symbol); const int loc = it->second.first; @@ -352,11 +356,9 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // parameter-shift we need to apply a single `gradient_gate` // per a symbol. std::complex result = ss.InnerProduct(scratch2, scratch); - auto val = - ( - std::complex(static_cast(coeff * result.real()), - static_cast(coeff * result.imag()))); - std::cout << ">>>>>>... value" << val << std::endl; + auto val = (std::complex( + static_cast(coeff * result.real()), + static_cast(coeff * result.imag()))); (*output_tensor)(i, loc, loc) += val; } } @@ -377,9 +379,7 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // other_sv contains sum_j other_programs_coeffs[i][j]*|phi[i][j]> // Start adjoint differentiation on two gates // m is the index for the first gate - std::cout << ">>> Start two gates hessian" << std::endl; for (int m = partial_fused_grad_circuits[i].size() - 1; m >= 1; m--) { - std::cout << ">>>>>>(1) m=" << m << "th partial fused circuit is applied" << std::endl; for (int k = partial_fused_grad_circuits[i][m].size() - 1; k >= 0; k--) { ApplyFusedGateDagger(sim, partial_fused_grad_circuits[i][m][k], sv); @@ -412,18 +412,15 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // non-controlled version of the gradient gate. ss.BulkSetAmpl(scratch4, mask_m, cbits_m, 0, 0, true); } - std::cout << ">>>>>>(1)... p=" << p << "th gradient gate is applied" << std::endl; qsim::ApplyGateDagger(sim, gradient_gates[i][m - 1].grad_gates[p], scratch4); // don't need not-found check since this is done upstream already. const auto it = maps[i].find(gradient_gates[i][m - 1].params[p]); - std::cout << ">>>>>>(1)... p=" << p << "th symbol = " << gradient_gates[i][m - 1].params[p] << std::endl; const int loc_m = it->second.first; // n is the index for the second gate for (int n = m - 1; n >= 0; n--) { - std::cout << ">>>>>>---(2) " << n << "th partial fused circuit is applied" << std::endl; for (int k = partial_fused_grad_circuits[i][n].size() - 1; k >= 0; k--) { ApplyFusedGateDagger(sim, partial_fused_grad_circuits[i][n][k], @@ -432,13 +429,11 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { scratch4); } if (n == 0) { - std::cout << ">>>>>>---(2) no just quit." << std::endl; // last layer will have no parametrized gates so can break. break; } // Hit a parameterized gate. - std::cout << "n-th gate index = " << gradient_gates[i][n - 1].index << std::endl; auto cur_gate_n = qsim_circuits[i].gates[gradient_gates[i][n - 1].index]; ApplyGateDagger(sim, cur_gate_n, scratch2); @@ -465,28 +460,11 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { ss.BulkSetAmpl(scratch3, mask_n, cbits_n, 0, 0, true); } - std::cout << ">>>>>>---(2)... q=" << q << "th gradient gate is applied" << std::endl; qsim::ApplyGate(sim, gradient_gates[i][n - 1].grad_gates[q], scratch3); - auto ptr = scratch3.get(); - auto ptr_size = 2 << scratch3.num_qubits(); - std::cout << "Statevector" << std::endl; - for (int i = 0; i < ptr_size; i++) { - std::cout << ptr[i] << ","; - } - std::cout << std::endl; - - ptr = scratch4.get(); - ptr_size = 2 << scratch4.num_qubits(); - std::cout << "Other Statevector" << std::endl; - for (int i = 0; i < ptr_size; i++) { - std::cout << ptr[i] << ","; - } - std::cout << std::endl; // don't need not-found check since this is done upstream already. const auto it = maps[i].find(gradient_gates[i][n - 1].params[q]); - std::cout << ">>>>>>---(2)... q=" << q << "th symbol = " << gradient_gates[i][n - 1].params[q] << std::endl; const int loc_n = it->second.first; // Apply finite differencing for adjoint gradients. // Finite differencing enables applying multiple `gradient_gate` @@ -495,11 +473,9 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // per a symbol. double coeff = static_cast(programs_coeffs[i]); std::complex result = ss.InnerProduct(scratch3, scratch4); - auto val = - ( - std::complex(static_cast(coeff * result.real()), - static_cast(coeff * result.imag()))); - std::cout << ">>>>>>>>>... value" << val << std::endl; + auto val = (std::complex( + static_cast(coeff * result.real()), + static_cast(coeff * result.imag()))); (*output_tensor)(i, loc_m, loc_n) += val; (*output_tensor)(i, loc_n, loc_m) += val; } @@ -533,7 +509,7 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { const int output_dim_internal_size = other_fused_circuits[0].size(); - auto DoWork = [&](int start, int end) { + auto DoWork1 = [&](int start, int end) { int old_batch_index = -2; int cur_batch_index = -1; int largest_nq = 1; @@ -542,13 +518,12 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { Simulator sim = Simulator(tfq_for); StateSpace ss = StateSpace(tfq_for); auto sv = ss.Create(largest_nq); + auto sv_adj = ss.Create(largest_nq); auto scratch = ss.Create(largest_nq); auto scratch2 = ss.Create(largest_nq); - auto scratch3 = ss.Create(largest_nq); - auto scratch4 = ss.Create(largest_nq); - for (int i = start; i < end; i++) { - cur_batch_index = i / output_dim_internal_size; - cur_internal_index = i % output_dim_internal_size; + for (int ii = start; ii < end; ii++) { + cur_batch_index = ii / output_dim_internal_size; + cur_internal_index = ii % output_dim_internal_size; const int nq = num_qubits[cur_batch_index]; @@ -558,6 +533,7 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { if (nq > largest_nq) { largest_nq = nq; sv = ss.Create(largest_nq); + sv_adj = ss.Create(largest_nq); scratch = ss.Create(largest_nq); scratch2 = ss.Create(largest_nq); } @@ -581,12 +557,18 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // now sv is |psi> // scratch contains |phi> // Start adjoint differentiation on a single gate - for (int l = partial_fused_hess_circuits[cur_batch_index].size() - 1; l >= 0; l--) { - for (int k = partial_fused_hess_circuits[cur_batch_index][l].size() - 1; k >= 0; - k--) { - ApplyFusedGateDagger(sim, partial_fused_hess_circuits[cur_batch_index][l][k], sv); - ApplyFusedGateDagger(sim, partial_fused_hess_circuits[cur_batch_index][l][k], - scratch); + ss.Copy(sv, sv_adj); + for (int l = partial_fused_hess_circuits[cur_batch_index].size() - 1; + l >= 0; l--) { + for (int k = + partial_fused_hess_circuits[cur_batch_index][l].size() - 1; + k >= 0; k--) { + ApplyFusedGateDagger( + sim, partial_fused_hess_circuits[cur_batch_index][l][k], + sv_adj); + ApplyFusedGateDagger( + sim, partial_fused_hess_circuits[cur_batch_index][l][k], + scratch); } if (l == 0) { // last layer will have no parametrized gates so can break. @@ -594,8 +576,10 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { } // Hit a parameterized gate. - auto cur_gate = qsim_circuits[cur_batch_index].gates[hessian_gates[cur_batch_index][l - 1].index]; - ApplyGateDagger(sim, cur_gate, sv); + auto cur_gate = + qsim_circuits[cur_batch_index] + .gates[hessian_gates[cur_batch_index][l - 1].index]; + ApplyGateDagger(sim, cur_gate, sv_adj); // if applicable compute control qubit mask and control value bits. uint64_t mask = 0; @@ -608,25 +592,34 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { } for (std::vector::size_type k = 0; - k < hessian_gates[cur_batch_index][l - 1].grad_gates.size(); k++) { + k < hessian_gates[cur_batch_index][l - 1].grad_gates.size(); + k++) { // Copy sv onto scratch2 in anticipation of non-unitary "gradient // gate". - ss.Copy(sv, scratch2); + ss.Copy(sv_adj, scratch2); if (!cur_gate.controlled_by.empty()) { // Gradient of controlled gates puts zeros on diagonal which is // the same as collapsing the state and then applying the // non-controlled version of the gradient gate. ss.BulkSetAmpl(scratch2, mask, cbits, 0, 0, true); } - qsim::ApplyGate(sim, hessian_gates[cur_batch_index][l - 1].grad_gates[k], scratch2); + qsim::ApplyGate(sim, + hessian_gates[cur_batch_index][l - 1].grad_gates[k], + scratch2); // don't need not-found check since this is done upstream already. auto symbol = hessian_gates[cur_batch_index][l - 1].params[k]; + double coeff1 = + static_cast(programs_coeffs[cur_batch_index]); + double coeff2 = static_cast( + other_programs_coeffs[cur_batch_index][cur_internal_index]); if (symbol == kUsePrevTwoSymbols) { // Apply second-order finite difference w.r.t. two symbols // That is, CrossTerm w.r.t. two symbols in one gate. - auto symbol1 = hessian_gates[cur_batch_index][l - 1].params[k - 2]; - auto symbol2 = hessian_gates[cur_batch_index][l - 1].params[k - 1]; + auto symbol1 = + hessian_gates[cur_batch_index][l - 1].params[k - 2]; + auto symbol2 = + hessian_gates[cur_batch_index][l - 1].params[k - 1]; auto it = maps[cur_batch_index].find(symbol1); const int loc1 = it->second.first; it = maps[cur_batch_index].find(symbol2); @@ -637,9 +630,9 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // parameter-shift we need to apply a single `gradient_gate` // per a symbol. std::complex result = ss.InnerProduct(scratch2, scratch); - auto val = (programs_coeffs[cur_batch_index] * other_programs_coeffs[cur_batch_index][cur_internal_index] * - std::complex(static_cast(result.real()), - static_cast(result.imag()))); + auto val = (std::complex( + static_cast(coeff1 * coeff2 * result.real()), + static_cast(coeff1 * coeff2 * result.imag()))); // Because Hessian is symmetric. (*output_tensor)(cur_batch_index, loc1, loc2) += val; (*output_tensor)(cur_batch_index, loc2, loc1) += val; @@ -654,35 +647,94 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // per a symbol. std::complex result = ss.InnerProduct(scratch2, scratch); (*output_tensor)(cur_batch_index, loc, loc) += - (programs_coeffs[cur_batch_index] * other_programs_coeffs[cur_batch_index][cur_internal_index] * - std::complex(static_cast(result.real()), - static_cast(result.imag()))); + (std::complex( + static_cast(coeff1 * coeff2 * result.real()), + static_cast(coeff1 * coeff2 * result.imag()))); } } ApplyGateDagger(sim, cur_gate, scratch); } + old_batch_index = cur_batch_index; + } + }; - // Re-initialize statevectors to save memory. - ss.SetStateZero(sv); - for (std::vector>::size_type j = 0; - j < fused_circuits[cur_batch_index].size(); j++) { - qsim::ApplyFusedGate(sim, fused_circuits[cur_batch_index][j], sv); + int64_t num_cycles = + 200 * (int64_t(1) << static_cast(max_num_qubits)); + context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor( + fused_circuits.size() * output_dim_internal_size, num_cycles, DoWork1); + + auto DoWork2 = [&](int start, int end) { + int old_batch_index = -2; + int cur_batch_index = -1; + int largest_nq = 1; + int cur_internal_index; + + Simulator sim = Simulator(tfq_for); + StateSpace ss = StateSpace(tfq_for); + auto sv = ss.Create(largest_nq); + auto sv_adj = ss.Create(largest_nq); + auto scratch = ss.Create(largest_nq); + auto scratch2 = ss.Create(largest_nq); + auto scratch3 = ss.Create(largest_nq); + auto scratch4 = ss.Create(largest_nq); + for (int ii = start; ii < end; ii++) { + cur_batch_index = ii / output_dim_internal_size; + cur_internal_index = ii % output_dim_internal_size; + + const int nq = num_qubits[cur_batch_index]; + + 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); + sv_adj = ss.Create(largest_nq); + scratch = ss.Create(largest_nq); + scratch2 = ss.Create(largest_nq); + scratch3 = ss.Create(largest_nq); + scratch4 = ss.Create(largest_nq); + } + ss.SetStateZero(sv); + for (std::vector>::size_type j = 0; + j < fused_circuits[cur_batch_index].size(); j++) { + qsim::ApplyFusedGate(sim, fused_circuits[cur_batch_index][j], sv); + } } + ss.SetStateZero(scratch); + for (std::vector>::size_type k = 0; + k < + other_fused_circuits[cur_batch_index][cur_internal_index].size(); + k++) { + qsim::ApplyFusedGate( + sim, other_fused_circuits[cur_batch_index][cur_internal_index][k], + scratch); + } + + // Re-initialize statevectors to save memory. + ss.Copy(sv, sv_adj); + // now sv is |psi> // scratch contains |phi> // Start adjoint differentiation on two gates // m is the index for the first gate - for (int m = partial_fused_grad_circuits[cur_batch_index].size() - 1; m >= 1; m--) { - for (int k = partial_fused_grad_circuits[cur_batch_index][m].size() - 1; k >= 0; - k--) { - ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][m][k], sv); - ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][m][k], - scratch); + for (int m = partial_fused_grad_circuits[cur_batch_index].size() - 1; + m >= 1; m--) { + for (int k = + partial_fused_grad_circuits[cur_batch_index][m].size() - 1; + k >= 0; k--) { + ApplyFusedGateDagger( + sim, partial_fused_grad_circuits[cur_batch_index][m][k], + sv_adj); + ApplyFusedGateDagger( + sim, partial_fused_grad_circuits[cur_batch_index][m][k], + scratch); } auto cur_gate_m = - qsim_circuits[cur_batch_index].gates[gradient_gates[cur_batch_index][m - 1].index]; - ApplyGateDagger(sim, cur_gate_m, sv); + qsim_circuits[cur_batch_index] + .gates[gradient_gates[cur_batch_index][m - 1].index]; + ApplyGateDagger(sim, cur_gate_m, sv_adj); // if applicable compute control qubit mask and control value bits. uint64_t mask_m = 0; @@ -693,34 +745,41 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { mask_m |= uint64_t{1} << control_loc; cbits_m |= ((cur_gate_m.cmask >> k) & 1) << control_loc; } + + ss.Copy(scratch, scratch4); + ss.Copy(sv_adj, scratch2); for (std::vector::size_type p = 0; - p < gradient_gates[cur_batch_index][m - 1].grad_gates.size(); p++) { + p < gradient_gates[cur_batch_index][m - 1].grad_gates.size(); + p++) { // Copy sv onto scratch2 in anticipation of the first non-unitary // "gradient gate". - ss.Copy(sv, scratch2); if (!cur_gate_m.controlled_by.empty()) { // Gradient of controlled gates puts zeros on diagonal which is // the same as collapsing the state and then applying the // non-controlled version of the gradient gate. - ss.BulkSetAmpl(scratch2, mask_m, cbits_m, 0, 0, true); + ss.BulkSetAmpl(scratch4, mask_m, cbits_m, 0, 0, true); } - qsim::ApplyGate(sim, gradient_gates[cur_batch_index][m - 1].grad_gates[p], - scratch2); + qsim::ApplyGateDagger( + sim, gradient_gates[cur_batch_index][m - 1].grad_gates[p], + scratch4); // don't need not-found check since this is done upstream already. - const auto it = maps[cur_batch_index].find(gradient_gates[cur_batch_index][m - 1].params[p]); + const auto it = maps[cur_batch_index].find( + gradient_gates[cur_batch_index][m - 1].params[p]); const int loc_m = it->second.first; - // Copy scratch onto scratch4. - ss.Copy(scratch, scratch4); // n is the index for the second gate for (int n = m - 1; n >= 0; n--) { - for (int k = partial_fused_grad_circuits[cur_batch_index][n].size() - 1; k >= 0; - k--) { - ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][n][k], - scratch2); - ApplyFusedGateDagger(sim, partial_fused_grad_circuits[cur_batch_index][n][k], - scratch4); + for (int k = + partial_fused_grad_circuits[cur_batch_index][n].size() - + 1; + k >= 0; k--) { + ApplyFusedGateDagger( + sim, partial_fused_grad_circuits[cur_batch_index][n][k], + scratch2); + ApplyFusedGateDagger( + sim, partial_fused_grad_circuits[cur_batch_index][n][k], + scratch4); } if (n == 0) { // last layer will have no parametrized gates so can break. @@ -729,10 +788,12 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // Hit a parameterized gate. auto cur_gate_n = - qsim_circuits[cur_batch_index].gates[gradient_gates[cur_batch_index][n - 1].index]; + qsim_circuits[cur_batch_index] + .gates[gradient_gates[cur_batch_index][n - 1].index]; ApplyGateDagger(sim, cur_gate_n, scratch2); - // if applicable compute control qubit mask and control value bits. + // if applicable compute control qubit mask and control value + // bits. uint64_t mask_n = 0; uint64_t cbits_n = 0; for (std::vector::size_type k = 0; @@ -743,36 +804,43 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { } for (std::vector::size_type q = 0; - q < gradient_gates[cur_batch_index][n - 1].grad_gates.size(); q++) { + q < gradient_gates[cur_batch_index][n - 1].grad_gates.size(); + q++) { // Copy scratch2 onto scratch3 in anticipation of the second // non-unitary "gradient gate". ss.Copy(scratch2, scratch3); if (!cur_gate_n.controlled_by.empty()) { - // Gradient of controlled gates puts zeros on diagonal which is - // the same as collapsing the state and then applying the + // Gradient of controlled gates puts zeros on diagonal which + // is the same as collapsing the state and then applying the // non-controlled version of the gradient gate. ss.BulkSetAmpl(scratch3, mask_n, cbits_n, 0, 0, true); } - qsim::ApplyGate(sim, gradient_gates[cur_batch_index][n - 1].grad_gates[q], - scratch3); - // don't need not-found check since this is done upstream already. - const auto it = maps[cur_batch_index].find(gradient_gates[cur_batch_index][n - 1].params[q]); + qsim::ApplyGate( + sim, gradient_gates[cur_batch_index][n - 1].grad_gates[q], + scratch3); + + // don't need not-found check since this is done upstream + // already. + const auto it = maps[cur_batch_index].find( + gradient_gates[cur_batch_index][n - 1].params[q]); const int loc_n = it->second.first; // Apply finite differencing for adjoint gradients. // Finite differencing enables applying multiple `gradient_gate` // of a symbol at the same circuit. For analytic methods like // parameter-shift we need to apply a single `gradient_gate` // per a symbol. - std::complex result = ss.InnerProduct(scratch3, scratch4); - auto val = - (programs_coeffs[cur_batch_index] * other_programs_coeffs[cur_batch_index][cur_internal_index] * - std::complex(static_cast(result.real()), - static_cast(result.imag()))); + double coeff1 = + static_cast(programs_coeffs[cur_batch_index]); + double coeff2 = static_cast( + other_programs_coeffs[cur_batch_index][cur_internal_index]); + std::complex result = + ss.InnerProduct(scratch3, scratch4); + auto val = (std::complex( + static_cast(coeff1 * coeff2 * result.real()), + static_cast(coeff1 * coeff2 * result.imag()))); (*output_tensor)(cur_batch_index, loc_m, loc_n) += val; - if (loc_m != loc_n) { - (*output_tensor)(cur_batch_index, loc_n, loc_m) += val; - } + (*output_tensor)(cur_batch_index, loc_n, loc_m) += val; } ApplyGateDagger(sim, cur_gate_n, scratch4); } @@ -783,10 +851,9 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { } }; - const int64_t num_cycles = - 200 * (int64_t(1) << static_cast(max_num_qubits)); + num_cycles = 500 * (int64_t(1) << static_cast(max_num_qubits)); context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor( - fused_circuits.size() * output_dim_internal_size, num_cycles, DoWork); + fused_circuits.size() * output_dim_internal_size, num_cycles, DoWork2); } }; diff --git a/tensorflow_quantum/core/src/adj_hessian_util.cc b/tensorflow_quantum/core/src/adj_hessian_util.cc index eb3970fcf..25be2c72f 100644 --- a/tensorflow_quantum/core/src/adj_hessian_util.cc +++ b/tensorflow_quantum/core/src/adj_hessian_util.cc @@ -36,7 +36,8 @@ void CreateHessianCircuit( const QsimCircuit& circuit, const std::vector& metadata, std::vector>>* partial_fuses, std::vector* grad_gates) { - for (std::vector::size_type i = 0; i < metadata.size(); i++) { + for (std::vector::size_type i = 0; i < metadata.size(); + i++) { if (metadata[i].symbol_values.empty()) { continue; } @@ -73,37 +74,10 @@ void CreateHessianCircuit( grad_gates->push_back(grad); } - // PhasedX + // Due to the large error from the 2nd order finite differencing + // PhasedXPowGate, we rejected the circuit if it contains the gate. else if (circuit.gates[i].kind == qsim::Cirq::GateKind::kPhasedXPowGate) { - // Process potentially several symbols. - bool symbolic_pexp = false; - bool symbolic_exp = false; - for (std::vector >::size_type j = 0; j < metadata[i].symbol_values.size(); j++) { - if (metadata[i].placeholder_names[j] == - GateParamNames::kPhaseExponent) { - symbolic_pexp = true; - PopulateHessianPhasedXPhasedExponent( - metadata[i].symbol_values[j], i, circuit.gates[i].qubits[0], - metadata[i].gate_params[0], metadata[i].gate_params[1], - metadata[i].gate_params[2], metadata[i].gate_params[3], - metadata[i].gate_params[4], &grad); - } else if (metadata[i].placeholder_names[j] == - GateParamNames::kExponent) { - symbolic_exp = true; - PopulateHessianPhasedXExponent( - metadata[i].symbol_values[j], i, circuit.gates[i].qubits[0], - metadata[i].gate_params[0], metadata[i].gate_params[1], - metadata[i].gate_params[2], metadata[i].gate_params[3], - metadata[i].gate_params[4], &grad); - } - } - if (symbolic_pexp && symbolic_exp) { - PopulateCrossTermPhasedXPhasedExponentExponent( - i, circuit.gates[i].qubits[0], metadata[i].gate_params[0], - metadata[i].gate_params[1], metadata[i].gate_params[2], - metadata[i].gate_params[3], metadata[i].gate_params[4], &grad); - } - grad_gates->push_back(grad); + // TODO(jaeyoo) : Add PhasedXPowGate hessian terms w/ small error. } // Fsim @@ -113,7 +87,8 @@ void CreateHessianCircuit( bool swapq = circuit.gates[i].swapped; bool symbolic_theta = false; bool symbolic_phi = false; - for (std::vector >::size_type j = 0; j < metadata[i].symbol_values.size(); j++) { + for (std::vector>::size_type j = 0; + j < metadata[i].symbol_values.size(); j++) { if (metadata[i].placeholder_names[j] == GateParamNames::kTheta) { symbolic_theta = true; PopulateHessianFsimTheta( @@ -150,7 +125,8 @@ void CreateHessianCircuit( bool swapq = circuit.gates[i].swapped; bool symbolic_pexp = false; bool symbolic_exp = false; - for (std::vector >::size_type j = 0; j < metadata[i].symbol_values.size(); j++) { + for (std::vector>::size_type j = 0; + j < metadata[i].symbol_values.size(); j++) { if (metadata[i].placeholder_names[j] == GateParamNames::kPhaseExponent) { symbolic_pexp = true; @@ -190,7 +166,8 @@ void CreateHessianCircuit( partial_fuses->assign(grad_gates->size() + 1, std::vector>({})); - for (std::vector::size_type i = 0; i < grad_gates->size(); i++) { + for (std::vector::size_type i = 0; i < grad_gates->size(); + i++) { right = circuit.gates.begin() + (*grad_gates)[i].index; (*partial_fuses)[i] = fuser.FuseGates(qsim::BasicGateFuser::Parameter(), @@ -218,77 +195,11 @@ void PopulateHessianSingleEigen( qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); - std::cout << "left = ["; - int size = 8; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = left.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; - std::cout << "right = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = right.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; - std::cout << "center = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = center.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; Matrix2Add(right.matrix, left.matrix); // left's entries have right added - std::cout << "left_plus_right = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = left.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto prefix = (val > 0) ? ((i % 2 == 1) ? "+" : "") : ""; - auto ending = (i % 2 == 0) ? "" : ((i == size - 1) ? "j" : "j,"); - std::cout << prefix << buf << ending; - } - std::cout << "]" << std::endl; qsim::MatrixScalarMultiply(2.0, center.matrix); - std::cout << "twice_center = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = center.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; Matrix2Diff(center.matrix, left.matrix); // left's entries have center subtracted. - std::cout << "left_plus_right_minus_twice_center = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = left.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; grad->grad_gates.push_back(left); } @@ -306,161 +217,11 @@ void PopulateHessianTwoEigen( qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); - std::cout << "PopulateHessianTwoEigen" << std::endl; - int size = 32; - std::cout << "left = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = left.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; - std::cout << "center = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = center.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; - std::cout << "right = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = right.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; Matrix4Add(right.matrix, left.matrix); // left's entries have right added. - std::cout << "left_plus_right = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = left.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; qsim::MatrixScalarMultiply(2.0, center.matrix); - std::cout << "twice_center = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = center.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; Matrix4Diff(center.matrix, left.matrix); // left's entries have center subtracted. - std::cout << "left_plus_right_minus_twice_center = ["; - for (int i = 0; i < size; i++) { - char buf[100]; - auto val = left.matrix[i]; - std::sprintf(buf, "%.7e", val); - auto ending = ( - (i % 2 == 0) ? ((val > 0) ? "+" : "") : - ((i == size - 1) ? "j" : "j,")); - std::cout << buf << ending; - } - std::cout << "]" << std::endl; - grad->grad_gates.push_back(left); -} - -void PopulateHessianPhasedXPhasedExponent(const std::string& symbol, - unsigned int location, - unsigned int qid, float pexp, - float pexp_s, float exp, float exp_s, - float gs, GradientOfGate* grad) { - grad->params.push_back(symbol); - grad->index = location; -// auto left = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, (pexp + _HESS_EPS) * pexp_s, exp * exp_s, gs); -// auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, -// exp * exp_s, gs); -// auto right = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, (pexp - _HESS_EPS) * pexp_s, exp * exp_s, gs); -// // Due to precision issue, multiply weights first. -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); -// Matrix2Add(right.matrix, -// left.matrix); // left's entries have right added. -// qsim::MatrixScalarMultiply(2.0, center.matrix); -// Matrix2Diff(center.matrix, -// left.matrix); // left's entries have center subtracted. - auto left = D2PhasedExponentPhasedXPowGate::Create( - 0, qid, pexp, pexp_s, exp*exp_s, gs); - grad->grad_gates.push_back(left); -} - -void PopulateHessianPhasedXExponent(const std::string& symbol, - unsigned int location, unsigned int qid, - float pexp, float pexp_s, float exp, - float exp_s, float gs, - GradientOfGate* grad) { - grad->params.push_back(symbol); - grad->index = location; -// auto left = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, pexp * pexp_s, (exp + _HESS_EPS) * exp_s, gs); -// auto center = qsim::Cirq::PhasedXPowGate::Create(0, qid, pexp * pexp_s, -// exp * exp_s, gs); -// auto right = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, pexp * pexp_s, (exp - _HESS_EPS) * exp_s, gs); -// // Due to precision issue, multiply weights first. -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, center.matrix); -// Matrix2Add(right.matrix, -// left.matrix); // left's entries have right added. -// qsim::MatrixScalarMultiply(2.0, center.matrix); -// Matrix2Diff(center.matrix, -// left.matrix); // left's entries have center subtracted. - auto left = D2ExponentPhasedXPowGate::Create( - 0, qid, pexp * pexp_s, exp, exp_s, gs); - grad->grad_gates.push_back(left); -} - -void PopulateCrossTermPhasedXPhasedExponentExponent( - unsigned int location, unsigned int qid, float pexp, float pexp_s, - float exp, float exp_s, float gs, GradientOfGate* grad) { - grad->params.push_back(kUsePrevTwoSymbols); - grad->index = location; -// auto left = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); -// auto left_center = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); -// auto right_center = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs); -// auto right = qsim::Cirq::PhasedXPowGate::Create( -// 0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs); -// // Due to precision issue, multiply weights first. -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left.matrix); -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right.matrix); -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, left_center.matrix); -// qsim::MatrixScalarMultiply(_INVERSE_HESS_EPS_SQUARE, right_center.matrix); -// Matrix2Add(right.matrix, -// left.matrix); // left's entries have right added. -// Matrix2Add(right_center.matrix, left_center.matrix); -// Matrix2Diff(left_center.matrix, -// left.matrix); // left's entries have left_center subtracted. - auto left = DPhasedExponentDExponentPhasedXPowGate::Create( - 0, qid, pexp, pexp_s, exp, exp_s, gs); grad->grad_gates.push_back(left); } diff --git a/tensorflow_quantum/core/src/adj_hessian_util.h b/tensorflow_quantum/core/src/adj_hessian_util.h index 4ad665f95..be8dae664 100644 --- a/tensorflow_quantum/core/src/adj_hessian_util.h +++ b/tensorflow_quantum/core/src/adj_hessian_util.h @@ -58,17 +58,6 @@ void PopulateHessianTwoEigen( // Note: all methods below expect gate qubit indices to have been swapped so // qid < qid2. -void PopulateHessianPhasedXPhasedExponent(const std::string& symbol, - unsigned int location, - unsigned int qid, float pexp, - float pexp_s, float exp, float exp_s, - float gs, GradientOfGate* grad); - -void PopulateHessianPhasedXExponent(const std::string& symbol, - unsigned int location, unsigned int qid, - float pexp, float pexp_s, float exp, - float exp_s, float gs, - GradientOfGate* grad); void PopulateCrossTermPhasedXPhasedExponentExponent( unsigned int location, unsigned int qid, float pexp, float pexp_s, @@ -119,141 +108,6 @@ void Matrix4Add(Array2& source, Array2& dest) { } } -// Due to the large error from the finite differencing PhasedXPowGate, -// Here we introduce analytically differentiated version of PhasedXPowGate up to -// 2nd order with double precision. - -template -struct D2PhasedExponentPhasedXPowGate { - static constexpr qsim::Cirq::GateKind kind = qsim::Cirq::kPhasedXPowGate; - static constexpr char name[] = "D2PhasedExponentPhasedXPowGate"; - static constexpr unsigned num_qubits = 1; - static constexpr bool symmetric = true; - - static qsim::Cirq::GateCirq Create(unsigned time, unsigned q0, - fp_type phase_exponent, - fp_type phase_exponent_scalar, - fp_type exponent = 1, - fp_type global_shift = 0) { - double pi = qsim::Cirq::pi_double; - double pexp = static_cast(phase_exponent); - double pexp_s = static_cast(phase_exponent_scalar); - double exp = static_cast(exponent); - double ec = std::cos(pi * exp); - double es = std::sin(pi * exp); - double gc = std::cos(pi * exp * global_shift); - double gs = std::sin(pi * exp * global_shift); - - double d2p_pc = -pi*pi*pexp_s*pexp_s*std::cos(pi * pexp_s * pexp); - double d2p_ps = -pi*pi*pexp_s*pexp_s*std::sin(pi * pexp_s * pexp); - - double br = -0.5 * ((-1 + ec) * gc - es * gs); - double bi = -0.5 * ((-1 + ec) * gs + es * gc); - - return qsim::CreateGate, D2PhasedExponentPhasedXPowGate>( - time, {q0}, {0., 0., static_cast(d2p_pc * br + d2p_ps * bi), - static_cast(d2p_pc * bi - d2p_ps * br), - static_cast(d2p_pc * br - d2p_ps * bi), - static_cast(d2p_pc * bi + d2p_ps * br), 0., 0.}, - {phase_exponent, exponent, global_shift}); - } -}; - -template -struct D2ExponentPhasedXPowGate { - static constexpr qsim::Cirq::GateKind kind = qsim::Cirq::kPhasedXPowGate; - static constexpr char name[] = "D2ExponentPhasedXPowGate"; - static constexpr unsigned num_qubits = 1; - static constexpr bool symmetric = true; - - static constexpr fp_type pi = static_cast(qsim::Cirq::pi_double); - - static qsim::Cirq::GateCirq Create(unsigned time, unsigned q0, - fp_type phase_exponent, - fp_type exponent = 1, - fp_type exponent_scalar = 1, - fp_type global_shift = 0) { - double pi = qsim::Cirq::pi_double; - double pexp = static_cast(phase_exponent); - double exp = static_cast(exponent); - double exp_s = static_cast(exponent_scalar); - double pc = std::cos(pi * pexp); - double ps = std::sin(pi * pexp); - double ec = std::cos(pi * exp_s * exp); - double es = std::sin(pi * exp_s * exp); - double gc = std::cos(pi * exp_s * exp * global_shift); - double gs = std::sin(pi * exp_s * exp * global_shift); - double dec = -pi * exp_s * std::sin(pi * exp_s * exp); - double des = pi * exp_s * std::cos(pi * exp_s * exp); - double dgc = -pi * exp_s * global_shift * std::sin(pi * exp_s * exp * global_shift); - double dgs = pi * exp_s * global_shift * std::cos(pi * exp_s * exp * global_shift); - double d2ec = - pi * exp_s * pi * exp_s * std::cos(pi * exp_s * exp); - double d2es = - pi * exp_s * pi * exp_s * std::sin(pi * exp_s * exp); - double d2gc = - pi * exp_s * pi * exp_s * global_shift * global_shift * std::cos(pi * exp_s * exp * global_shift); - double d2gs = - pi * exp_s * pi * exp_s * global_shift * global_shift * std::sin(pi * exp_s * exp * global_shift); - double common_r_front = ec * d2gc + 2.0 * dec * dgc + d2ec * gc; - double common_r_back = d2es * gs + 2.0 * des * dgs + es * d2gs; - double common_i_front = ec * d2gs + 2.0 * dec * dgs + d2ec * gs; - double common_i_back = d2es * gc + 2.0 * des * dgc + es * d2gc; - - fp_type d2ar = static_cast(0.5 * (d2gc + common_r_front - common_r_back)); - fp_type d2ai = static_cast(0.5 * (d2gs + common_i_front + common_i_back)); - double d2br = -0.5 * (-d2gc + common_r_front - common_r_back); - double d2bi = -0.5 * (-d2gs + common_i_front + common_i_back); - - return qsim::CreateGate, D2ExponentPhasedXPowGate>( - time, {q0}, {d2ar, d2ai, static_cast(pc * d2br + ps * d2bi), - static_cast(pc * d2bi - ps * d2br), - static_cast(pc * d2br - ps * d2bi), - static_cast(pc * d2bi + ps * d2br), d2ar, d2ai}, - {phase_exponent, exponent, global_shift}); - } -}; - -template -struct DPhasedExponentDExponentPhasedXPowGate { - static constexpr qsim::Cirq::GateKind kind = qsim::Cirq::kPhasedXPowGate; - static constexpr char name[] = "DPhasedExponentDExponentPhasedXPowGate"; - static constexpr unsigned num_qubits = 1; - static constexpr bool symmetric = true; - - static constexpr fp_type pi = static_cast(qsim::Cirq::pi_double); - - static qsim::Cirq::GateCirq Create(unsigned time, unsigned q0, - fp_type phase_exponent, - fp_type phase_exponent_scalar, - fp_type exponent = 1, - fp_type exponent_scalar = 1, - fp_type global_shift = 0) { - double pi = qsim::Cirq::pi_double; - double pexp = static_cast(phase_exponent); - double pexp_s = static_cast(phase_exponent_scalar); - double exp = static_cast(exponent); - double exp_s = static_cast(exponent_scalar); - double ec = std::cos(pi * exp_s * exp); - double es = std::sin(pi * exp_s * exp); - double gc = std::cos(pi * exp_s * exp * global_shift); - double gs = std::sin(pi * exp_s * exp * global_shift); - - double dp_pc = -pi * pexp_s * std::sin(pi * pexp_s * pexp); - double dp_ps = pi * pexp_s * std::cos(pi * pexp_s * pexp); - double de_ec = -pi * exp_s * std::sin(pi * exp_s * exp); - double de_es = pi * exp_s * std::cos(pi * exp_s * exp); - double de_gc = -pi * exp_s * global_shift * std::sin(pi * exp_s * global_shift * exp); - double de_gs = pi * exp_s * global_shift * std::cos(pi * exp_s * global_shift * exp); - - double de_br = -0.5 * ((-1 + ec) * de_gc + de_ec * gc - de_es * gs - es * de_gs); - double de_bi = -0.5 * ((-1 + ec) * de_gs + de_ec * gs + de_es * gc + es * de_gc); - - return qsim::CreateGate, DPhasedExponentDExponentPhasedXPowGate>( - time, {q0}, {0.,0., static_cast(dp_pc * de_br + dp_ps * de_bi), - static_cast(dp_pc * de_bi - dp_ps * de_br), - static_cast(dp_pc * de_br - dp_ps * de_bi), - static_cast(dp_pc * de_bi + dp_ps * de_br), 0.,0.}, - {phase_exponent, exponent, global_shift}); - } -}; - } // namespace tfq #endif // TFQ_CORE_SRC_ADJ_UTIL_H_ diff --git a/tensorflow_quantum/python/util.py b/tensorflow_quantum/python/util.py index 72a1df558..bd9ac8e5a 100644 --- a/tensorflow_quantum/python/util.py +++ b/tensorflow_quantum/python/util.py @@ -39,7 +39,7 @@ ] -def get_supported_gates(): +def get_supported_gates(exclude_gates=None): """A helper to get gates supported by TFQ. Returns a dictionary mapping from supported gate types @@ -50,8 +50,10 @@ def get_supported_gates(): supported. """ supported_ops = serializer.SERIALIZER.supported_gate_types() - supported_gates = filter(lambda x: x not in _SUPPORTED_CHANNELS, - supported_ops) + exclude_gates = exclude_gates if exclude_gates else [] + supported_gates = filter( + lambda x: not (x in _SUPPORTED_CHANNELS or x in exclude_gates), + supported_ops) gate_arity_mapping_dict = dict() for gate in supported_gates: if gate is cirq.IdentityGate: @@ -123,14 +125,15 @@ def random_symbol_circuit(qubits, n_moments=15, p=0.9, include_scalars=True, - include_channels=False): + include_channels=False, + exclude_gates=None): """Generates a random circuit including some parameterized gates. Symbols are randomly included in the gates of the first `n_moments` moments of the resulting circuit. Then, parameterized H gates are added as subsequent moments for any remaining unused symbols. """ - supported_ops = get_supported_gates() + supported_ops = get_supported_gates(exclude_gates) if include_channels: for chan, n in get_supported_channels().items(): supported_ops[chan] = n @@ -217,7 +220,8 @@ def random_symbol_circuit_resolver_batch(qubits, n_moments=15, p=0.9, include_scalars=True, - include_channels=False): + include_channels=False, + exclude_gates=None): """Generate a batch of random circuits and resolvers.""" return_circuits = [] return_resolvers = [] @@ -228,7 +232,8 @@ def random_symbol_circuit_resolver_batch(qubits, n_moments=n_moments, p=p, include_scalars=include_scalars, - include_channels=include_channels)) + include_channels=include_channels, + exclude_gates=exclude_gates)) return_resolvers.append( cirq.ParamResolver( From 1c3f25c52a7b535358bf6fafd13aa26ac1e2a1f2 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 09:28:52 +0900 Subject: [PATCH 04/12] Add inner_product_hessian import_test --- scripts/import_test.py | 1 + .../core/ops/math_ops/__init__.py | 1 + .../math_ops/inner_product_hessian_test.py | 428 +++++++++--------- 3 files changed, 218 insertions(+), 212 deletions(-) diff --git a/scripts/import_test.py b/scripts/import_test.py index a92a09795..1cb2e69f8 100644 --- a/scripts/import_test.py +++ b/scripts/import_test.py @@ -36,6 +36,7 @@ def test_imports(): # Math ops. _ = tfq.math.inner_product + _ = tfq.math.inner_product_hessian # Noisy simulation ops. _ = tfq.noise.expectation diff --git a/tensorflow_quantum/core/ops/math_ops/__init__.py b/tensorflow_quantum/core/ops/math_ops/__init__.py index 7842b8784..72a4e637b 100644 --- a/tensorflow_quantum/core/ops/math_ops/__init__.py +++ b/tensorflow_quantum/core/ops/math_ops/__init__.py @@ -15,3 +15,4 @@ """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.inner_product_op import inner_product_hessian diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py index 707ad1fa9..b22e9899c 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -52,6 +52,7 @@ def get_gate(gate, symbol_names, qubits): + """Generates a gate operation.""" symbols = sympy.symbols(symbol_names) if len(symbols) == 1: a, b = symbols * 2 @@ -75,6 +76,7 @@ def get_gate(gate, symbol_names, qubits): def get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, resolver): + """Generates a state vector with shifted values.""" new_resolver = copy.deepcopy(resolver) new_resolver.param_dict[name_j] += dx_j new_resolver.param_dict[name_k] += dx_k @@ -82,6 +84,7 @@ def get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, resolver): def get_finite_difference_hessian(circuit, name_j, name_k, resolver): + """Generates finite difference hessian.""" # dx came from _GRAD_EPS of core/src/adj_util.cc dx = 5e-3 inv_square_two_dx = np.asarray([1e4 + 0.j], dtype=np.complex64) @@ -103,209 +106,209 @@ def get_finite_difference_hessian(circuit, name_j, name_k, resolver): class InnerProductAdjHessianTest(tf.test.TestCase, parameterized.TestCase): - """Tests tfq_inner_product_hessian.""" - # - # def test_inner_product_hessian_inputs(self): - # """Makes sure that inner_product_adj_hessian fails on bad inputs.""" - # n_qubits = 5 - # batch_size = 5 - # n_other_programs = 3 - # symbol_names = ['alpha'] - # qubits = cirq.GridQubit.rect(1, n_qubits) - # programs_coeffs = np.ones((batch_size,)) - # other_programs_coeffs = np.ones((batch_size, n_other_programs)) - # circuit_batch, resolver_batch = \ - # util.random_symbol_circuit_resolver_batch( - # qubits, symbol_names, batch_size) - # - # symbol_values_array = np.array( - # [[resolver[symbol] - # for symbol in symbol_names] - # for resolver in resolver_batch]) - # - # other_batch = [ - # util.random_circuit_resolver_batch(qubits, n_other_programs)[0] - # for _ in range(batch_size) - # ] - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'programs must be rank 1'): - # # Circuit tensor has too many dimensions. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor([circuit_batch]), - # symbol_names, symbol_values_array, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'symbol_names must be rank 1.'): - # # symbol_names tensor has too many dimensions. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), - # np.array([symbol_names]), symbol_values_array, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'symbol_values must be rank 2.'): - # # symbol_values_array tensor has too many dimensions. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # np.array([symbol_values_array]), - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'symbol_values must be rank 2.'): - # # symbol_values_array tensor has too few dimensions. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), - # symbol_names, symbol_values_array[0], - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'other_programs must be rank 2.'): - # # other_programs tensor has too few dimensions. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), - # symbol_names, symbol_values_array, - # util.convert_to_tensor(circuit_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'other_programs must be rank 2.'): - # # pauli_sums tensor has too many dimensions. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array, - # util.convert_to_tensor([[x] for x in other_batch]), - # programs_coeffs, other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'Unparseable proto'): - # # circuit tensor has the right type but invalid values. - # inner_product_op.inner_product_hessian( - # ['junk'] * batch_size, symbol_names, symbol_values_array, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'Could not find symbol in parameter map'): - # # symbol_names tensor has the right type but invalid values. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), - # ['junk'], symbol_values_array, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'not found in reference circuit'): - # # other_programs tensor has the right type but operates on - # # qubits that the reference ciruit doesn't have. - # new_qubits = [cirq.GridQubit(5, 5), cirq.GridQubit(9, 9)] - # new_circuits, _ = util.random_circuit_resolver_batch( - # new_qubits, batch_size) - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array, - # util.convert_to_tensor([[x] for x in new_circuits]), - # programs_coeffs, other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # 'not found in paired circuit'): - # # other_programs tensor has the right type but operates on - # # qubits that the reference ciruit doesn't have. - # new_qubits = cirq.GridQubit.rect(1, n_qubits - 1) - # new_circuits, _ = util.random_circuit_resolver_batch( - # new_qubits, batch_size) - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array, - # util.convert_to_tensor([[x] for x in new_circuits]), - # programs_coeffs, other_programs_coeffs) - # - # with self.assertRaisesRegex(TypeError, 'Cannot convert'): - # # circuits tensor has the wrong type. - # inner_product_op.inner_product_hessian( - # [1.0] * batch_size, symbol_names, symbol_values_array, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(TypeError, 'Cannot convert'): - # # symbol_names tensor has the wrong type. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), - # [0.1234], symbol_values_array, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.UnimplementedError, ''): - # # symbol_values tensor has the wrong type. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), - # symbol_names, [['junk']] * batch_size, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(TypeError, 'Cannot convert'): - # # other_programs tensor has the wrong type. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array, [[1.0]] * batch_size, programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex(TypeError, 'missing'): - # # we are missing an argument. - # # pylint: disable=no-value-for-parameter - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array, programs_coeffs, other_programs_coeffs) - # # pylint: enable=no-value-for-parameter - # - # with self.assertRaisesRegex(TypeError, 'positional arguments'): - # # pylint: disable=too-many-function-args - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), - # symbol_names, symbol_values_array, - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs, []) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # expected_regex='do not match'): - # # batch programs has wrong batch size. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array, - # util.convert_to_tensor(other_batch[:int(batch_size * 0.5)]), - # programs_coeffs, other_programs_coeffs) - # - # with self.assertRaisesRegex(tf.errors.InvalidArgumentError, - # expected_regex='do not match'): - # # batch programs has wrong batch size. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array[::int(batch_size * 0.5)], - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # - # with self.assertRaisesRegex( - # tf.errors.InvalidArgumentError, - # expected_regex='Found symbols in other_programs'): - # # other_programs has symbols. - # inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array, - # util.convert_to_tensor([[x] for x in circuit_batch]), - # programs_coeffs, other_programs_coeffs) - # - # res = inner_product_op.inner_product_hessian( - # util.convert_to_tensor(circuit_batch), symbol_names, - # symbol_values_array.astype(np.float64), - # util.convert_to_tensor(other_batch), programs_coeffs, - # other_programs_coeffs) - # self.assertDTypeEqual(res, np.complex64) - # + """Tests tfq_inner_product_hessian."""s + + def test_inner_product_hessian_inputs(self): + """Makes sure that inner_product_adj_hessian fails on bad inputs.""" + n_qubits = 5 + batch_size = 5 + n_other_programs = 3 + symbol_names = ['alpha'] + qubits = cirq.GridQubit.rect(1, n_qubits) + programs_coeffs = np.ones((batch_size,)) + other_programs_coeffs = np.ones((batch_size, n_other_programs)) + circuit_batch, resolver_batch = \ + util.random_symbol_circuit_resolver_batch( + qubits, symbol_names, batch_size) + + symbol_values_array = np.array( + [[resolver[symbol] + for symbol in symbol_names] + for resolver in resolver_batch]) + + other_batch = [ + util.random_circuit_resolver_batch(qubits, n_other_programs)[0] + for _ in range(batch_size) + ] + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'programs must be rank 1'): + # Circuit tensor has too many dimensions. + inner_product_op.inner_product_hessian( + util.convert_to_tensor([circuit_batch]), + symbol_names, symbol_values_array, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbol_names must be rank 1.'): + # symbol_names tensor has too many dimensions. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), + np.array([symbol_names]), symbol_values_array, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbol_values must be rank 2.'): + # symbol_values_array tensor has too many dimensions. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + np.array([symbol_values_array]), + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'symbol_values must be rank 2.'): + # symbol_values_array tensor has too few dimensions. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), + symbol_names, symbol_values_array[0], + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'other_programs must be rank 2.'): + # other_programs tensor has too few dimensions. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), + symbol_names, symbol_values_array, + util.convert_to_tensor(circuit_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'other_programs must be rank 2.'): + # pauli_sums tensor has too many dimensions. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in other_batch]), + programs_coeffs, other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'Unparseable proto'): + # circuit tensor has the right type but invalid values. + inner_product_op.inner_product_hessian( + ['junk'] * batch_size, symbol_names, symbol_values_array, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'Could not find symbol in parameter map'): + # symbol_names tensor has the right type but invalid values. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), + ['junk'], symbol_values_array, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'not found in reference circuit'): + # other_programs tensor has the right type but operates on + # qubits that the reference ciruit doesn't have. + new_qubits = [cirq.GridQubit(5, 5), cirq.GridQubit(9, 9)] + new_circuits, _ = util.random_circuit_resolver_batch( + new_qubits, batch_size) + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in new_circuits]), + programs_coeffs, other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + 'not found in paired circuit'): + # other_programs tensor has the right type but operates on + # qubits that the reference ciruit doesn't have. + new_qubits = cirq.GridQubit.rect(1, n_qubits - 1) + new_circuits, _ = util.random_circuit_resolver_batch( + new_qubits, batch_size) + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in new_circuits]), + programs_coeffs, other_programs_coeffs) + + with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # circuits tensor has the wrong type. + inner_product_op.inner_product_hessian( + [1.0] * batch_size, symbol_names, symbol_values_array, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # symbol_names tensor has the wrong type. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), + [0.1234], symbol_values_array, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.UnimplementedError, ''): + # symbol_values tensor has the wrong type. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), + symbol_names, [['junk']] * batch_size, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(TypeError, 'Cannot convert'): + # other_programs tensor has the wrong type. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, [[1.0]] * batch_size, programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex(TypeError, 'missing'): + # we are missing an argument. + # pylint: disable=no-value-for-parameter + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, programs_coeffs, other_programs_coeffs) + # pylint: enable=no-value-for-parameter + + with self.assertRaisesRegex(TypeError, 'positional arguments'): + # pylint: disable=too-many-function-args + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), + symbol_names, symbol_values_array, + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs, []) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='do not match'): + # batch programs has wrong batch size. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor(other_batch[:int(batch_size * 0.5)]), + programs_coeffs, other_programs_coeffs) + + with self.assertRaisesRegex(tf.errors.InvalidArgumentError, + expected_regex='do not match'): + # batch programs has wrong batch size. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array[::int(batch_size * 0.5)], + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + expected_regex='Found symbols in other_programs'): + # other_programs has symbols. + inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array, + util.convert_to_tensor([[x] for x in circuit_batch]), + programs_coeffs, other_programs_coeffs) + + res = inner_product_op.inner_product_hessian( + util.convert_to_tensor(circuit_batch), symbol_names, + symbol_values_array.astype(np.float64), + util.convert_to_tensor(other_batch), programs_coeffs, + other_programs_coeffs) + self.assertDTypeEqual(res, np.complex64) + @parameterized.parameters([ { 'n_qubits': 5, @@ -330,13 +333,14 @@ class InnerProductAdjHessianTest(tf.test.TestCase, parameterized.TestCase): ]) def test_correctness_with_symbols(self, n_qubits, batch_size, inner_dim_size): - """Tests that inner_product works with symbols.""" + """Tests that inner_product_hessian works with symbols.""" symbol_names = ['alpha', 'beta', 'gamma', 'delta', 'eta', 'kappa'] n_params = len(symbol_names) qubits = cirq.GridQubit.rect(1, n_qubits) circuit_batch, resolver_batch = \ util.random_symbol_circuit_resolver_batch( - qubits, symbol_names, batch_size, exclude_gates=_UNSUPPORTED_GATES) + qubits, symbol_names, batch_size, + exclude_gates=_UNSUPPORTED_GATES) other_batch = [ util.random_circuit_resolver_batch(qubits, inner_dim_size)[0] @@ -410,7 +414,7 @@ def test_correctness_with_symbols(self, n_qubits, batch_size, }, ]) def correctness_without_symbols(self, n_qubits, batch_size, inner_dim_size): - """Tests that inner_product_adj_grad works without symbols.""" + """Tests that inner_product_hessian works without symbols.""" qubits = cirq.GridQubit.rect(1, n_qubits) circuit_batch, _ = \ util.random_circuit_resolver_batch( @@ -437,7 +441,7 @@ def correctness_without_symbols(self, n_qubits, batch_size, inner_dim_size): other_programs_coeffs) def correctness_empty(self): - """Tests the inner product adj grad between two empty circuits.""" + """Tests the inner product hessian between two empty circuits.""" symbol_names = ['alpha', 'beta'] n_params = len(symbol_names) empty_cicuit = util.convert_to_tensor([cirq.Circuit()]) @@ -469,8 +473,7 @@ def correctness_empty(self): self.assertAllClose(out, expected) def correctness_no_circuit(self): - """Test the inner product grad between no circuits.""" - + """Test the inner product hessian between no 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) @@ -489,6 +492,7 @@ def correctness_no_circuit(self): class InnerProductHessianOnGates(tf.test.TestCase, parameterized.TestCase): + """Tests inner_product_hessian on a single gate.""" @parameterized.parameters([{ 'gate': gate, @@ -496,8 +500,8 @@ class InnerProductHessianOnGates(tf.test.TestCase, parameterized.TestCase): } for gate in _ONE_EIGEN_GATES + _TWO_EIGEN_GATES for names in _SYMBOL_NAMES]) - def correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): - """Tests that inner_product works with symbols.""" + def correctness_one_gate_with_symbols(self, gate, symbol_names): + """Tests that inner_product_hessian works with one gate.""" n_params = len(symbol_names) qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] @@ -548,7 +552,7 @@ def correctness_one_qubit_gate_with_symbols(self, gate, symbol_names): 'gate': gate, } for gate in _UNSUPPORTED_GATES for names in _SYMBOL_NAMES]) def unsupported_gate_with_symbols(self, gate): - """Tests that inner_product works with symbols.""" + """Tests that inner_product_hessian deals with unsupported gates.""" symbol_names = ['alpha'] qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] From 2e44e0348b09a6438e3cb01c756b8baf8a8a0b12 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 09:31:33 +0900 Subject: [PATCH 05/12] Fix lint --- tensorflow_quantum/core/ops/math_ops/__init__.py | 4 ++-- .../core/ops/math_ops/inner_product_hessian_test.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow_quantum/core/ops/math_ops/__init__.py b/tensorflow_quantum/core/ops/math_ops/__init__.py index 72a4e637b..8358e07ec 100644 --- a/tensorflow_quantum/core/ops/math_ops/__init__.py +++ b/tensorflow_quantum/core/ops/math_ops/__init__.py @@ -14,5 +14,5 @@ # ============================================================================== """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.inner_product_op import inner_product_hessian +from tensorflow_quantum.core.ops.math_ops.inner_product_op import ( + inner_product, inner_product_hessian) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py index b22e9899c..259309cbf 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -106,7 +106,7 @@ def get_finite_difference_hessian(circuit, name_j, name_k, resolver): class InnerProductAdjHessianTest(tf.test.TestCase, parameterized.TestCase): - """Tests tfq_inner_product_hessian."""s + """Tests tfq_inner_product_hessian.""" def test_inner_product_hessian_inputs(self): """Makes sure that inner_product_adj_hessian fails on bad inputs.""" From 2a1259494361ef64f34912924ff891512aaa299e Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 09:32:31 +0900 Subject: [PATCH 06/12] Fix format --- tensorflow_quantum/core/ops/math_ops/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow_quantum/core/ops/math_ops/__init__.py b/tensorflow_quantum/core/ops/math_ops/__init__.py index 8358e07ec..b7d6f0936 100644 --- a/tensorflow_quantum/core/ops/math_ops/__init__.py +++ b/tensorflow_quantum/core/ops/math_ops/__init__.py @@ -15,4 +15,4 @@ """Module for tfq.core.ops.math_ops.*""" from tensorflow_quantum.core.ops.math_ops.inner_product_op import ( - inner_product, inner_product_hessian) + inner_product, inner_product_hessian) From 67a16f900bcc3571b728494bf8c3cc2fcd4ba2f3 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 10:01:42 +0900 Subject: [PATCH 07/12] Fix format --- .../core/ops/math_ops/inner_product_op.py | 122 ++++++++++++------ .../core/src/adj_hessian_util.h | 7 +- tensorflow_quantum/core/src/adj_util.cc | 2 + tensorflow_quantum/core/src/adj_util.h | 2 - 4 files changed, 84 insertions(+), 49 deletions(-) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py index db82bd999..6d4d7e0ac 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py @@ -24,48 +24,86 @@ def inner_product_hessian(programs, symbol_names, symbol_values, other_programs, programs_coeffs, other_programs_coeffs): """Calculate the adjoint Hessian of the inner product between circuits. - Compute the gradients of the (potentially many) inner products between - the given circuits and the symbol free comparison circuits. - - Calculates out[i][j][k] = $\text{programs_coeffs[i]} \times \langle - \frac{\partial^2 \psi_{\text{programs[i]}}(\text{symbol_values[i]})} - {\partial \text{symbol_names[j]}\partial \text{symbol_names[k]}} | - \sum_l \text{other_programs_coeffs[l]}\times | - \psi_{\text{other_programs[l]}} \rangle$ - - - Note: `other_programs` must not contain any free symbols. These can - be resolved beforehand with `tfq.resolve_parameters`. - - Note: len(symbol_names) (=n_params) should be a positive integer. - - 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. - programs_coeffs: `tf.Tensor` of real numbers with shape [batch_size] - of weights on `programs`. - other_programs_coeffs: `tf.Tensor` of real numbers with shape - [batch_size, n_others] of weights on `other_programs`. - - Returns: - tf.Tensor` with shape [batch_size, n_symbols, n_symbols] where `out[i]` is - equal to the hessian of the inner product between programs[i] and all - other_programs[i] w.r.t. `symbol_names[j]` and `symbol_names[k]`. - `programs[i]` is resolved with `symbol_values[i]` and each - (other_)programs[i] is weighted by (other_)programs_coeffs[i]. - """ + Compute the gradients of the (potentially many) inner products between + the given circuits and the symbol free comparison circuits. + + Calculates out[i][j][k] = $\text{programs_coeffs[i]} \times \langle + \frac{\partial^2 \psi_{\text{programs[i]}}(\text{symbol_values[i]})} + {\partial \text{symbol_names[j]}\partial \text{symbol_names[k]}} | + \sum_l \text{other_programs_coeffs[l]}\times | + \psi_{\text{other_programs[l]}} \rangle$ + + + >>> symbols = sympy.symbols('alpha beta') + >>> qubits = cirq.GridQubit.rect(1, 2) + >>> reference_circuits = [ + ... cirq.Circuit((cirq.H**symbols[0]).on_each(qubits)), + ... cirq.Circuit( + ... cirq.X(qubits[0]) ** symbols[0], + ... cirq.Y(qubits[1]) ** symbols[1]) + ... ] + >>> other_circuits = [ + ... cirq.Circuit(cirq.X.on_each(qubits)), + ... cirq.Circuit((cirq.Y**0.125).on_each(qubits)), + ... cirq.Circuit((cirq.X**0.5).on_each(qubits)) + ... ] + >>> reference_tensor = tfq.convert_to_tensor(reference_circuits) + >>> symbol_tensor = tf.convert_to_tensor([s.name for s in symbols]) + >>> values_tensor = tf.convert_to_tensor(np.arange(4).reshape(2, 2)) + >>> other_tensor = tfq.convert_to_tensor([other_circuits, other_circuits]) + >>> reference_coeff_tensor = tf.ones((len(reference_circuits,))) + >>> other_coeff_tensor = tf.ones((len(reference_circuits), + ... len(other_circuits))) + >>> ip = tfq.math.inner_product_hessian(reference_tensor, symbol_tensor, + ... values_tensor, other_tensor, + ... reference_coeff_tensor, + ... other_coeff_tensor) + >>> ip + tf.Tensor( + [[[ 0.6069082-1.0185852j 0. +0.j ] + [ 0. +0.j 0. +0.j ]] + + [[-2.7567296-1.7676406j -0.8554697+1.0770144j] + [-0.8554697+1.0770144j 4.0239005+7.6238985j]]], shape=(2, 2, 2), + dtype=complex64) + + + + Note: `other_programs` must not contain any free symbols. These can + be resolved beforehand with `tfq.resolve_parameters`. + + Note: `programs` must not contain `cirq.PhasedXPowGate` due to precision + issue. + + Note: len(symbol_names) (=n_params) should be a positive integer. + + 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. + programs_coeffs: `tf.Tensor` of real numbers with shape [batch_size] + of weights on `programs`. + other_programs_coeffs: `tf.Tensor` of real numbers with shape + [batch_size, n_others] of weights on `other_programs`. + + Returns: + tf.Tensor` with shape [batch_size, n_params, n_params] where `out[i]` + is equal to the hessian of the inner product between programs[i] and all + other_programs[i] w.r.t. `symbol_names[j]` and `symbol_names[k]`. + `programs[i]` is resolved with `symbol_values[i]` and each + (other_)programs[i] is weighted by (other_)programs_coeffs[i]. + """ return MATH_OP_MODULE.tfq_inner_product_hessian( programs, symbol_names, tf.cast(symbol_values, tf.float32), other_programs, tf.cast(programs_coeffs, tf.float32), diff --git a/tensorflow_quantum/core/src/adj_hessian_util.h b/tensorflow_quantum/core/src/adj_hessian_util.h index be8dae664..d589caaae 100644 --- a/tensorflow_quantum/core/src/adj_hessian_util.h +++ b/tensorflow_quantum/core/src/adj_hessian_util.h @@ -31,6 +31,7 @@ limitations under the License. namespace tfq { +static const float _GRAD_EPS = 5e-3; static const float _HESS_EPS = 1e-2; static const float _INVERSE_HESS_EPS_SQUARE = 1e4; static const std::string kUsePrevTwoSymbols = "use_prev_two_symbols"; @@ -58,11 +59,6 @@ void PopulateHessianTwoEigen( // Note: all methods below expect gate qubit indices to have been swapped so // qid < qid2. - -void PopulateCrossTermPhasedXPhasedExponentExponent( - unsigned int location, unsigned int qid, float pexp, float pexp_s, - float exp, float exp_s, float gs, GradientOfGate* grad); - void PopulateHessianFsimTheta(const std::string& symbol, unsigned int location, unsigned int qid, unsigned qid2, float theta, float theta_s, float phi, float phi_s, @@ -93,6 +89,7 @@ void PopulateCrossTermPhasedISwapPhasedExponentExponent( unsigned int location, unsigned int qid, unsigned int qid2, float pexp, float pexp_s, float exp, float exp_s, GradientOfGate* grad); +// does matrix elementwise addition dest += source. template void Matrix2Add(Array2& source, Array2& dest) { for (unsigned i = 0; i < 8; i++) { diff --git a/tensorflow_quantum/core/src/adj_util.cc b/tensorflow_quantum/core/src/adj_util.cc index c60bf7d57..ceb76b2c1 100644 --- a/tensorflow_quantum/core/src/adj_util.cc +++ b/tensorflow_quantum/core/src/adj_util.cc @@ -29,6 +29,8 @@ limitations under the License. namespace tfq { +static const float _GRAD_EPS = 5e-3; + typedef qsim::Cirq::GateCirq QsimGate; typedef qsim::Circuit QsimCircuit; diff --git a/tensorflow_quantum/core/src/adj_util.h b/tensorflow_quantum/core/src/adj_util.h index b9913f1de..b4bd30385 100644 --- a/tensorflow_quantum/core/src/adj_util.h +++ b/tensorflow_quantum/core/src/adj_util.h @@ -30,8 +30,6 @@ limitations under the License. namespace tfq { -static const float _GRAD_EPS = 5e-3; - struct GradientOfGate { // name of parameters used by gate. std::vector params; From 01412620ae822053a65948bbeed501dfa0da1b86 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 10:03:48 +0900 Subject: [PATCH 08/12] Fix timeout by reducing the number of symbols --- .../core/ops/math_ops/inner_product_hessian_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py index 259309cbf..969918a5f 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -334,7 +334,7 @@ def test_inner_product_hessian_inputs(self): def test_correctness_with_symbols(self, n_qubits, batch_size, inner_dim_size): """Tests that inner_product_hessian works with symbols.""" - symbol_names = ['alpha', 'beta', 'gamma', 'delta', 'eta', 'kappa'] + symbol_names = ['alpha', 'beta', 'gamma'] n_params = len(symbol_names) qubits = cirq.GridQubit.rect(1, n_qubits) circuit_batch, resolver_batch = \ From 5f6ffcdd92e6fd80c7a759576732b93ba7ca1824 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 10:15:57 +0900 Subject: [PATCH 09/12] Fix test inputs --- .../math_ops/inner_product_hessian_test.py | 93 ++++--------------- tensorflow_quantum/core/ops/parse_context.cc | 2 +- 2 files changed, 20 insertions(+), 75 deletions(-) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py index 969918a5f..1650da052 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -51,7 +51,7 @@ ] -def get_gate(gate, symbol_names, qubits): +def _get_gate(gate, symbol_names, qubits): """Generates a gate operation.""" symbols = sympy.symbols(symbol_names) if len(symbols) == 1: @@ -75,7 +75,8 @@ def get_gate(gate, symbol_names, qubits): ] -def get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, resolver): +def _get_shifted_resolved_circuit(circuit, name_j, name_k, dx_j, dx_k, + resolver): """Generates a state vector with shifted values.""" new_resolver = copy.deepcopy(resolver) new_resolver.param_dict[name_j] += dx_j @@ -88,14 +89,14 @@ def get_finite_difference_hessian(circuit, name_j, name_k, resolver): # dx came from _GRAD_EPS of core/src/adj_util.cc dx = 5e-3 inv_square_two_dx = np.asarray([1e4 + 0.j], dtype=np.complex64) - final_circuit_pp = get_shifted_resolved_circuit(circuit, name_j, name_k, dx, - dx, resolver) - final_circuit_mp = get_shifted_resolved_circuit(circuit, name_j, name_k, - -dx, dx, resolver) - final_circuit_pm = get_shifted_resolved_circuit(circuit, name_j, name_k, dx, - -dx, resolver) - final_circuit_mm = get_shifted_resolved_circuit(circuit, name_j, name_k, - -dx, -dx, resolver) + final_circuit_pp = _get_shifted_resolved_circuit(circuit, name_j, name_k, + dx, dx, resolver) + final_circuit_mp = _get_shifted_resolved_circuit(circuit, name_j, name_k, + -dx, dx, resolver) + final_circuit_pm = _get_shifted_resolved_circuit(circuit, name_j, name_k, + dx, -dx, resolver) + final_circuit_mm = _get_shifted_resolved_circuit(circuit, name_j, name_k, + -dx, -dx, resolver) final_wf_pp = inv_square_two_dx * cirq.final_state_vector(final_circuit_pp) final_wf_mp = inv_square_two_dx * cirq.final_state_vector(final_circuit_mp) final_wf_pm = inv_square_two_dx * cirq.final_state_vector(final_circuit_pm) @@ -119,7 +120,8 @@ def test_inner_product_hessian_inputs(self): other_programs_coeffs = np.ones((batch_size, n_other_programs)) circuit_batch, resolver_batch = \ util.random_symbol_circuit_resolver_batch( - qubits, symbol_names, batch_size) + qubits, symbol_names, batch_size, + exclude_gates=_UNSUPPORTED_GATES) symbol_values_array = np.array( [[resolver[symbol] @@ -413,7 +415,8 @@ def test_correctness_with_symbols(self, n_qubits, batch_size, 'inner_dim_size': 5 }, ]) - def correctness_without_symbols(self, n_qubits, batch_size, inner_dim_size): + def test_correctness_without_symbols(self, n_qubits, batch_size, + inner_dim_size): """Tests that inner_product_hessian works without symbols.""" qubits = cirq.GridQubit.rect(1, n_qubits) circuit_batch, _ = \ @@ -440,7 +443,7 @@ def correctness_without_symbols(self, n_qubits, batch_size, inner_dim_size): progams_coeffs, other_programs_coeffs) - def correctness_empty(self): + def test_correctness_empty(self): """Tests the inner product hessian between two empty circuits.""" symbol_names = ['alpha', 'beta'] n_params = len(symbol_names) @@ -472,7 +475,7 @@ def correctness_empty(self): expected = np.zeros((1, n_params, n_params), dtype=np.complex64) self.assertAllClose(out, expected) - def correctness_no_circuit(self): + def test_correctness_no_circuit(self): """Test the inner product hessian between no circuits.""" empty_circuit = tf.raw_ops.Empty(shape=(0,), dtype=tf.string) empty_symbols = tf.raw_ops.Empty(shape=(0,), dtype=tf.string) @@ -490,72 +493,14 @@ def correctness_no_circuit(self): empty_circuit, empty_symbols, empty_values, other_program, empty_program_coeffs, empty_other_program_coeffs) - -class InnerProductHessianOnGates(tf.test.TestCase, parameterized.TestCase): - """Tests inner_product_hessian on a single gate.""" - - @parameterized.parameters([{ - 'gate': gate, - 'symbol_names': names - } - for gate in _ONE_EIGEN_GATES + _TWO_EIGEN_GATES - for names in _SYMBOL_NAMES]) - def correctness_one_gate_with_symbols(self, gate, symbol_names): - """Tests that inner_product_hessian works with one gate.""" - n_params = len(symbol_names) - qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) - circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] - resolver_batch = [ - cirq.ParamResolver({name: 0.123 for name in symbol_names}) - ] - - symbol_values_array = np.array( - [[resolver[symbol] - for symbol in symbol_names] - for resolver in resolver_batch]) - other_batch = [ - [cirq.Circuit(cirq.H.on_each(*qubits))] for _ in circuit_batch - ] - programs = util.convert_to_tensor(circuit_batch) - other_programs = util.convert_to_tensor(other_batch) - symbol_names_tensor = tf.convert_to_tensor(symbol_names, - dtype=tf.dtypes.string) - symbol_values = tf.convert_to_tensor(symbol_values_array) - programs_coeffs = tf.cast(tf.random.normal((1,)), tf.complex64) - other_programs_coeffs = tf.cast(tf.random.normal((1, 1)), tf.complex64) - - out = inner_product_op.inner_product_hessian( - programs, symbol_names_tensor, symbol_values, other_programs, - programs_coeffs, other_programs_coeffs) - - out_arr = np.zeros((1, n_params, n_params), dtype=np.complex64) - for i, resolver in enumerate(resolver_batch): - weighted_internal_wf = None - for l, other in enumerate(other_batch[i]): - internal_wf = (other_programs_coeffs[i][l] * - cirq.final_state_vector(other)) - if l == 0: - weighted_internal_wf = internal_wf - else: - weighted_internal_wf += internal_wf - for j, name_j in enumerate(symbol_names): - for k, name_k in enumerate(symbol_names): - final_wf_grad = get_finite_difference_hessian( - circuit_batch[i], name_j, name_k, resolver) - out_arr[i][j][k] += ( - programs_coeffs[i] * - np.vdot(final_wf_grad, weighted_internal_wf)) - - self.assertAllClose(out, out_arr, atol=_ATOL, rtol=_RTOL) - @parameterized.parameters([{ 'gate': gate, } for gate in _UNSUPPORTED_GATES for names in _SYMBOL_NAMES]) - def unsupported_gate_with_symbols(self, gate): + def test_unsupported_gate_with_symbols(self, gate): """Tests that inner_product_hessian deals with unsupported gates.""" symbol_names = ['alpha'] qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) - circuit_batch = [cirq.Circuit(get_gate(gate, symbol_names, qubits))] + circuit_batch = [cirq.Circuit(_get_gate(gate, symbol_names, qubits))] resolver_batch = [ cirq.ParamResolver({name: 0.123 for name in symbol_names}) ] diff --git a/tensorflow_quantum/core/ops/parse_context.cc b/tensorflow_quantum/core/ops/parse_context.cc index 55eb23003..59f5f3c70 100644 --- a/tensorflow_quantum/core/ops/parse_context.cc +++ b/tensorflow_quantum/core/ops/parse_context.cc @@ -217,7 +217,7 @@ tensorflow::Status GetProgramsAndNumQubits( if (programs->size() != other_programs->size()) { return Status(tensorflow::error::INVALID_ARGUMENT, absl::StrCat("programs and other_programs batch dimension", - " do not match. Foud: ", programs->size(), + " do not match. Found: ", programs->size(), " and ", other_programs->size())); } From 4e01d1d8f9eac2842c6e59ab52edc61de0550fe5 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 21:35:42 +0900 Subject: [PATCH 10/12] Fix name --- .../core/ops/math_ops/inner_product_hessian_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py index 1650da052..d7abdeaa2 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -29,13 +29,13 @@ _ATOL = 0.2 _RTOL = 0.3 _SYMBOL_NAMES = [['alpha'], ['alpha', 'beta']] -_ONE_EIGEN_GATES = [ +_ONE_QUBIT_GATES = [ cirq.XPowGate, cirq.YPowGate, cirq.ZPowGate, cirq.HPowGate, ] -_TWO_EIGEN_GATES = [ +_TWO_QUBIT_GATES = [ cirq.XXPowGate, cirq.YYPowGate, cirq.ZZPowGate, @@ -499,7 +499,7 @@ def test_correctness_no_circuit(self): def test_unsupported_gate_with_symbols(self, gate): """Tests that inner_product_hessian deals with unsupported gates.""" symbol_names = ['alpha'] - qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_EIGEN_GATES else 1) + qubits = cirq.GridQubit.rect(1, 2 if gate in _TWO_QUBIT_GATES else 1) circuit_batch = [cirq.Circuit(_get_gate(gate, symbol_names, qubits))] resolver_batch = [ cirq.ParamResolver({name: 0.123 for name in symbol_names}) From acea17b9b4b5888449b7983e7b10e1f35085604d Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 21:50:45 +0900 Subject: [PATCH 11/12] Fix typo --- tensorflow_quantum/core/ops/math_ops/inner_product_op.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py index 6d4d7e0ac..1608c8fce 100644 --- a/tensorflow_quantum/core/ops/math_ops/inner_product_op.py +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_op.py @@ -117,7 +117,7 @@ def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, Compute the gradients of the (potentially many) inner products between the given circuits and the symbol free comparison circuits. - Calculates out[i][j][k] = $\langle \frac{\partial \psi_{\text{programs[i]}} + Calculates out[i][j] = $\langle \frac{\partial \psi_{\text{programs[i]}} (\text{symbol_values[i]})}{\partial \text{symbol_names[j]}} | \sum_k \text{prev_grad[i][k]}\times | \psi_{\text{other_programs[k]}} \rangle$ @@ -131,7 +131,7 @@ def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, 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_symbols], which + 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`. @@ -147,7 +147,7 @@ def _inner_product_grad(programs, symbol_names, symbol_values, other_programs, backprop of values from downstream in the compute graph. Returns: - tf.Tensor` with shape [batch_size, n_symbols] where `out[i][j]` is equal + tf.Tensor` with shape [batch_size, n_params] where `out[i][j]` is equal to the gradient of the inner product between programs[i] and all other_programs[i] w.r.t. `symbol_names[j]` and `programs[i]` is resolved with `symbol_values[i]`. From 178910860f1c42c030c4c8fd0bf181e77c13cd90 Mon Sep 17 00:00:00 2001 From: Jae Yoo Date: Sat, 3 Apr 2021 23:10:14 +0900 Subject: [PATCH 12/12] Add mutex in ComputeSmall --- .../core/ops/math_ops/tfq_inner_product_hessian.cc | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc index 0930d4d25..8ff0c152d 100644 --- a/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc +++ b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc @@ -509,6 +509,7 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { const int output_dim_internal_size = other_fused_circuits[0].size(); + auto c_lock = tensorflow::mutex(); auto DoWork1 = [&](int start, int end) { int old_batch_index = -2; int cur_batch_index = -1; @@ -634,8 +635,10 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { static_cast(coeff1 * coeff2 * result.real()), static_cast(coeff1 * coeff2 * result.imag()))); // Because Hessian is symmetric. + c_lock.lock(); (*output_tensor)(cur_batch_index, loc1, loc2) += val; (*output_tensor)(cur_batch_index, loc2, loc1) += val; + c_lock.unlock(); } else { // Apply second-order finite difference w.r.t. one symbol const auto it = maps[cur_batch_index].find(symbol); @@ -646,10 +649,12 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { // parameter-shift we need to apply a single `gradient_gate` // per a symbol. std::complex result = ss.InnerProduct(scratch2, scratch); + c_lock.lock(); (*output_tensor)(cur_batch_index, loc, loc) += (std::complex( static_cast(coeff1 * coeff2 * result.real()), static_cast(coeff1 * coeff2 * result.imag()))); + c_lock.unlock(); } } ApplyGateDagger(sim, cur_gate, scratch); @@ -839,8 +844,10 @@ class TfqInnerProductHessianOp : public tensorflow::OpKernel { auto val = (std::complex( static_cast(coeff1 * coeff2 * result.real()), static_cast(coeff1 * coeff2 * result.imag()))); + c_lock.lock(); (*output_tensor)(cur_batch_index, loc_m, loc_n) += val; (*output_tensor)(cur_batch_index, loc_n, loc_m) += val; + c_lock.unlock(); } ApplyGateDagger(sim, cur_gate_n, scratch4); }