Add analytic diff on PhasedXPowGate

Author: jaeyoo

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)

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<QsimGate>::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);
           }

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<float>::Create(
-      0, qid, (pexp + _HESS_EPS) * pexp_s, exp * exp_s, gs);
-  auto center = qsim::Cirq::PhasedXPowGate<float>::Create(0, qid, pexp * pexp_s,
-                                                          exp * exp_s, gs);
-  auto right = qsim::Cirq::PhasedXPowGate<float>::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<float>::Create(
+//      0, qid, (pexp + _HESS_EPS) * pexp_s, exp * exp_s, gs);
+//  auto center = qsim::Cirq::PhasedXPowGate<float>::Create(0, qid, pexp * pexp_s,
+//                                                          exp * exp_s, gs);
+//  auto right = qsim::Cirq::PhasedXPowGate<float>::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<float>::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<float>::Create(
-      0, qid, pexp * pexp_s, (exp + _HESS_EPS) * exp_s, gs);
-  auto center = qsim::Cirq::PhasedXPowGate<float>::Create(0, qid, pexp * pexp_s,
-                                                          exp * exp_s, gs);
-  auto right = qsim::Cirq::PhasedXPowGate<float>::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<float>::Create(
+//      0, qid, pexp * pexp_s, (exp + _HESS_EPS) * exp_s, gs);
+//  auto center = qsim::Cirq::PhasedXPowGate<float>::Create(0, qid, pexp * pexp_s,
+//                                                          exp * exp_s, gs);
+//  auto right = qsim::Cirq::PhasedXPowGate<float>::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<float>::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<float>::Create(
-      0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs);
-  auto left_center = qsim::Cirq::PhasedXPowGate<float>::Create(
-      0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs);
-  auto right_center = qsim::Cirq::PhasedXPowGate<float>::Create(
-      0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs);
-  auto right = qsim::Cirq::PhasedXPowGate<float>::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<float>::Create(
+//      0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs);
+//  auto left_center = qsim::Cirq::PhasedXPowGate<float>::Create(
+//      0, qid, (pexp + _GRAD_EPS) * pexp_s, (exp - _GRAD_EPS) * exp_s, gs);
+//  auto right_center = qsim::Cirq::PhasedXPowGate<float>::Create(
+//      0, qid, (pexp - _GRAD_EPS) * pexp_s, (exp + _GRAD_EPS) * exp_s, gs);
+//  auto right = qsim::Cirq::PhasedXPowGate<float>::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<float>::Create(
+      0, qid, pexp, pexp_s, exp, exp_s, gs);
   grad->grad_gates.push_back(left);
 }