diff --git a/scripts/import_test.py b/scripts/import_test.py index 10cd3aa5a..0b08847fe 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/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/__init__.py b/tensorflow_quantum/core/ops/math_ops/__init__.py index 7842b8784..b7d6f0936 100644 --- a/tensorflow_quantum/core/ops/math_ops/__init__.py +++ b/tensorflow_quantum/core/ops/math_ops/__init__.py @@ -14,4 +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, inner_product_hessian) 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..d7abdeaa2 --- /dev/null +++ b/tensorflow_quantum/core/ops/math_ops/inner_product_hessian_test.py @@ -0,0 +1,531 @@ +# 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 = 0.2 +_RTOL = 0.3 +_SYMBOL_NAMES = [['alpha'], ['alpha', 'beta']] +_ONE_QUBIT_GATES = [ + cirq.XPowGate, + cirq.YPowGate, + cirq.ZPowGate, + cirq.HPowGate, +] +_TWO_QUBIT_GATES = [ + cirq.XXPowGate, + cirq.YYPowGate, + cirq.ZZPowGate, + cirq.CZPowGate, + cirq.CNotPowGate, + cirq.SwapPowGate, + cirq.ISwapPowGate, + cirq.PhasedISwapPowGate, + cirq.FSimGate, +] +_UNSUPPORTED_GATES = [ + cirq.PhasedXPowGate, +] + + +def _get_gate(gate, symbol_names, qubits): + """Generates a gate operation.""" + 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): + """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 + return cirq.resolve_parameters(circuit, new_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) + 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, + exclude_gates=_UNSUPPORTED_GATES) + + 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 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'] + 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) + + 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, 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_hessian 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 hessian 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 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) + 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([{ + 'gate': gate, + } for gate in _UNSUPPORTED_GATES for names in _SYMBOL_NAMES]) + 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_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}) + ] + + 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__": + 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..1608c8fce 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,96 @@ 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$ + + + >>> 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), + 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 +117,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] = $\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 @@ -52,11 +143,11 @@ 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: - 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]`. 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..8ff0c152d --- /dev/null +++ b/tensorflow_quantum/core/ops/math_ops/tfq_inner_product_hessian.cc @@ -0,0 +1,910 @@ +/* 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 "../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); + + 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]); + 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 20 since Memory = 4GB + // e2s4 = 4 CPU, 16GB -> Can safely do 20 since Memory = 8GB + // ... + 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, + 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 + for (int l = partial_fused_hess_circuits[i].size() - 1; l >= 0; l--) { + 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); + } + 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]; + 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]; + 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. + (*output_tensor)(i, loc1, loc2) += val; + (*output_tensor)(i, loc2, loc1) += val; + } else { + // 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()))); + (*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> + // 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 + for (int m = partial_fused_grad_circuits[i].size() - 1; m >= 1; m--) { + 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; + } + + 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". + 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(scratch4, mask_m, cbits_m, 0, 0, true); + } + 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]); + const int loc_m = it->second.first; + + // n is the index for the second gate + for (int n = m - 1; n >= 0; n--) { + 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) { + // 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); + } + + 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]); + 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()))); + (*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 c_lock = tensorflow::mutex(); + auto DoWork1 = [&](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); + 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); + } + 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 + 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. + 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_adj); + + // 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_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); + + // 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 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 = (std::complex( + 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); + 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); + 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); + } + old_batch_index = cur_batch_index; + } + }; + + 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_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_adj); + + // 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; + } + + 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++) { + // Copy sv onto scratch2 in anticipation of the first non-unitary + // "gradient gate". + 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(scratch4, mask_m, cbits_m, 0, 0, true); + } + 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 int loc_m = it->second.first; + + // 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. + 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()))); + 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); + } + } + ApplyGateDagger(sim, cur_gate_m, scratch); + } + old_batch_index = cur_batch_index; + } + }; + + 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, DoWork2); + } +}; + +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..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())); } @@ -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..25be2c72f --- /dev/null +++ b/tensorflow_quantum/core/src/adj_hessian_util.cc @@ -0,0 +1,378 @@ +/* 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); + } + + // 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) { + // TODO(jaeyoo) : Add PhasedXPowGate hessian terms w/ small error. + } + + // 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); + 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 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); + 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 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..d589caaae --- /dev/null +++ b/tensorflow_quantum/core/src/adj_hessian_util.h @@ -0,0 +1,110 @@ +/* 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 _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"; + +// 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 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); + +// does matrix elementwise addition dest += source. +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.h b/tensorflow_quantum/core/src/adj_util.h index 7cc383ad0..b4bd30385 100644 --- a/tensorflow_quantum/core/src/adj_util.h +++ b/tensorflow_quantum/core/src/adj_util.h @@ -100,7 +100,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 +108,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++) { 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(