Fix type mismatch warnings

tensorflow_quantum/core/ops/math_ops/tfq_inner_product.cc

@@ -179,7 +179,7 @@ class TfqInnerProductOp : public tensorflow::OpKernel {
     // 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 (int i = 0; i < fused_circuits.size(); i++) {
+    for (size_t i = 0; i < fused_circuits.size(); i++) {
       int nq = num_qubits[i];
       if (nq > largest_nq) {
         // need to switch to larger statespace.
@@ -191,18 +191,18 @@ class TfqInnerProductOp : public tensorflow::OpKernel {
       //  the state if there is a possibility that circuit[i] and
       //  circuit[i + 1] produce the same state.
       ss.SetStateZero(sv);
-      for (int j = 0; j < fused_circuits[i].size(); j++) {
+      for (size_t j = 0; j < fused_circuits[i].size(); j++) {
         qsim::ApplyFusedGate(sim, fused_circuits[i][j], sv);
       }
-      for (int j = 0; j < other_fused_circuits[i].size(); j++) {
+      for (size_t j = 0; j < other_fused_circuits[i].size(); j++) {
         // (#679) Just ignore empty program
         if (fused_circuits[i].size() == 0) {
           (*output_tensor)(i, j) = std::complex<float>(1, 0);
           continue;
         }
 
         ss.SetStateZero(scratch);
-        for (int k = 0; k < other_fused_circuits[i][j].size(); k++) {
+        for (size_t k = 0; k < other_fused_circuits[i][j].size(); k++) {
           qsim::ApplyFusedGate(sim, other_fused_circuits[i][j][k], scratch);
         }
 
@@ -260,13 +260,13 @@ class TfqInnerProductOp : public tensorflow::OpKernel {
           // no need to update scratch_state since ComputeExpectation
           // will take care of things for us.
           ss.SetStateZero(sv);
-          for (int j = 0; j < fused_circuits[cur_batch_index].size(); j++) {
+          for (size_t j = 0; j < fused_circuits[cur_batch_index].size(); j++) {
             qsim::ApplyFusedGate(sim, fused_circuits[cur_batch_index][j], sv);
           }
         }
 
         ss.SetStateZero(scratch);
-        for (int k = 0;
+        for (size_t k = 0;
              k <
              other_fused_circuits[cur_batch_index][cur_internal_index].size();
              k++) {

tensorflow_quantum/core/ops/math_ops/tfq_inner_product_grad.cc

@@ -62,8 +62,8 @@ class TfqInnerProductGradOp : public tensorflow::OpKernel {
 
     // 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);
+    size_t output_dim_internal_size = context->input(3).dim_size(1);
+    size_t 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 ",
@@ -403,13 +403,13 @@ class TfqInnerProductGradOp : public tensorflow::OpKernel {
           // if applicable compute control qubit mask and control value bits.
           uint64_t mask = 0;
           uint64_t cbits = 0;
-          for (int k = 0; k < cur_gate.controlled_by.size(); k++) {
+          for (size_t 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 (int k = 0;
+          for (size_t k = 0;
                k < gradient_gates[cur_batch_index][l - 1].grad_gates.size();
                k++) {
             // Copy sv_adj onto scratch2 in anticipation of non-unitary

tensorflow_quantum/core/ops/noise/BUILD

@@ -60,6 +60,7 @@ cc_binary(
         "//tensorflow_quantum/core/ops:parse_context",
         "//tensorflow_quantum/core/ops:tfq_simulate_utils",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
+        "//tensorflow_quantum/core/src:util_balance_trajectory",
         "//tensorflow_quantum/core/src:util_qsim",
         "@qsim//lib:qsim_lib",
         # tensorflow core framework

tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {
@@ -180,8 +181,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
     tensorflow::GuardedPhiloxRandom random_gen;
     int max_n_shots = 1;
-    for (int i = 0; i < num_samples.size(); i++) {
-      for (int j = 0; j < num_samples[i].size(); j++) {
+    for (size_t i = 0; i < num_samples.size(); i++) {
+      for (size_t j = 0; j < num_samples[i].size(); j++) {
         max_n_shots = std::max(max_n_shots, num_samples[i][j]);
       }
     }
@@ -193,12 +194,12 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
     // 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 (int i = 0; i < ncircuits.size(); i++) {
+    for (size_t i = 0; i < ncircuits.size(); i++) {
       int nq = num_qubits[i];
 
       // (#679) Just ignore empty program
       if (ncircuits[i].channels.size() == 0) {
-        for (int j = 0; j < pauli_sums[i].size(); j++) {
+        for (size_t j = 0; j < pauli_sums[i].size(); j++) {
           (*output_tensor)(i, j) = -2.0;
         }
         continue;
@@ -225,7 +226,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
                              sv, unused_stats);
 
         // Use this trajectory as a source for all expectation calculations.
-        for (int j = 0; j < pauli_sums[i].size(); j++) {
+        for (size_t j = 0; j < pauli_sums[i].size(); j++) {
           if (run_samples[j] >= num_samples[i][j]) {
             continue;
           }
@@ -237,14 +238,14 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
           run_samples[j]++;
         }
         bool break_loop = true;
-        for (int j = 0; j < num_samples[i].size(); j++) {
+        for (size_t j = 0; j < num_samples[i].size(); j++) {
           if (run_samples[j] < num_samples[i][j]) {
             break_loop = false;
             break;
           }
         }
         if (break_loop) {
-          for (int j = 0; j < num_samples[i].size(); j++) {
+          for (size_t j = 0; j < num_samples[i].size(); j++) {
             rolling_sums[j] /= num_samples[i][j];
             (*output_tensor)(i, j) = static_cast<float>(rolling_sums[j]);
           }
@@ -285,8 +286,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
     tensorflow::GuardedPhiloxRandom random_gen;
     int max_n_shots = 1;
-    for (int i = 0; i < num_samples.size(); i++) {
-      for (int j = 0; j < num_samples[i].size(); j++) {
+    for (size_t i = 0; i < num_samples.size(); i++) {
+      for (size_t j = 0; j < num_samples[i].size(); j++) {
         max_n_shots = std::max(max_n_shots, num_samples[i][j]);
       }
     }
@@ -309,13 +310,13 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
           random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
 
-      for (int i = 0; i < ncircuits.size(); i++) {
+      for (size_t i = 0; i < ncircuits.size(); i++) {
         int nq = num_qubits[i];
         int rep_offset = rep_offsets[start][i];
 
         // (#679) Just ignore empty program
         if (ncircuits[i].channels.size() == 0) {
-          for (int j = 0; j < pauli_sums[i].size(); j++) {
+          for (size_t j = 0; j < pauli_sums[i].size(); j++) {
             (*output_tensor)(i, j) = -2.0;
           }
           continue;
@@ -342,7 +343,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
                                sim, sv, unused_stats);
 
           // Compute expectations across all ops using this trajectory.
-          for (int j = 0; j < pauli_sums[i].size(); j++) {
+          for (size_t j = 0; j < pauli_sums[i].size(); j++) {
             int p_reps = (num_samples[i][j] + num_threads - 1) / num_threads;
             if (run_samples[j] >= p_reps + rep_offset) {
               continue;
@@ -359,7 +360,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
           // Check if we have run enough trajectories for all ops.
           bool break_loop = true;
-          for (int j = 0; j < num_samples[i].size(); j++) {
+          for (size_t j = 0; j < num_samples[i].size(); j++) {
             int p_reps = (num_samples[i][j] + num_threads - 1) / num_threads;
             if (run_samples[j] < p_reps + rep_offset) {
               break_loop = false;
@@ -369,7 +370,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
           if (break_loop) {
             // Lock writing to this batch index in output_tensor.
             batch_locks[i].lock();
-            for (int j = 0; j < num_samples[i].size(); j++) {
+            for (size_t j = 0; j < num_samples[i].size(); j++) {
               rolling_sums[j] /= num_samples[i][j];
               (*output_tensor)(i, j) += static_cast<float>(rolling_sums[j]);
             }

tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc

@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {
@@ -182,8 +183,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
     tensorflow::GuardedPhiloxRandom random_gen;
     int max_psum_length = 1;
     int max_n_shots = 1;
-    for (int i = 0; i < pauli_sums.size(); i++) {
-      for (int j = 0; j < pauli_sums[i].size(); j++) {
+    for (size_t i = 0; i < pauli_sums.size(); i++) {
+      for (size_t j = 0; j < pauli_sums[i].size(); j++) {
         max_psum_length =
             std::max(max_psum_length, pauli_sums[i][j].terms().size());
         max_n_shots = std::max(max_n_shots, num_samples[i][j]);
@@ -197,12 +198,12 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
     // 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 (int i = 0; i < ncircuits.size(); i++) {
+    for (size_t i = 0; i < ncircuits.size(); i++) {
       int nq = num_qubits[i];
 
       // (#679) Just ignore empty program
       if (ncircuits[i].channels.empty()) {
-        for (int j = 0; j < pauli_sums[i].size(); j++) {
+        for (size_t j = 0; j < pauli_sums[i].size(); j++) {
           (*output_tensor)(i, j) = -2.0;
         }
         continue;
@@ -229,7 +230,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
                              sv, unused_stats);
 
         // Use this trajectory as a source for all expectation calculations.
-        for (int j = 0; j < pauli_sums[i].size(); j++) {
+        for (size_t j = 0; j < pauli_sums[i].size(); j++) {
           if (run_samples[j] >= num_samples[i][j]) {
             continue;
           }
@@ -241,14 +242,14 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
           run_samples[j]++;
         }
         bool break_loop = true;
-        for (int j = 0; j < num_samples[i].size(); j++) {
+        for (size_t j = 0; j < num_samples[i].size(); j++) {
           if (run_samples[j] < num_samples[i][j]) {
             break_loop = false;
             break;
           }
         }
         if (break_loop) {
-          for (int j = 0; j < num_samples[i].size(); j++) {
+          for (size_t j = 0; j < num_samples[i].size(); j++) {
             rolling_sums[j] /= num_samples[i][j];
             (*output_tensor)(i, j) = static_cast<float>(rolling_sums[j]);
           }
@@ -290,8 +291,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
     tensorflow::GuardedPhiloxRandom random_gen;
     int max_psum_length = 1;
     int max_n_shots = 1;
-    for (int i = 0; i < pauli_sums.size(); i++) {
-      for (int j = 0; j < pauli_sums[i].size(); j++) {
+    for (size_t i = 0; i < pauli_sums.size(); i++) {
+      for (size_t j = 0; j < pauli_sums[i].size(); j++) {
         max_psum_length =
             std::max(max_psum_length, pauli_sums[i][j].terms().size());
         max_n_shots = std::max(max_n_shots, num_samples[i][j]);
@@ -315,13 +316,13 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
       auto local_gen = random_gen.ReserveSamples128(num_rand);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
 
-      for (int i = 0; i < ncircuits.size(); i++) {
+      for (size_t i = 0; i < ncircuits.size(); i++) {
         int nq = num_qubits[i];
         int rep_offset = rep_offsets[start][i];
 
         // (#679) Just ignore empty program
         if (ncircuits[i].channels.empty()) {
-          for (int j = 0; j < pauli_sums[i].size(); j++) {
+          for (size_t j = 0; j < pauli_sums[i].size(); j++) {
             (*output_tensor)(i, j) = -2.0;
           }
           continue;
@@ -348,7 +349,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
                                sim, sv, unused_stats);
 
           // Compute expectations across all ops using this trajectory.
-          for (int j = 0; j < pauli_sums[i].size(); j++) {
+          for (size_t j = 0; j < pauli_sums[i].size(); j++) {
             int p_reps = (num_samples[i][j] + num_threads - 1) / num_threads;
             if (run_samples[j] >= p_reps + rep_offset) {
               continue;
@@ -365,7 +366,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
 
           // Check if we have run enough trajectories for all ops.
           bool break_loop = true;
-          for (int j = 0; j < num_samples[i].size(); j++) {
+          for (size_t j = 0; j < num_samples[i].size(); j++) {
             int p_reps = (num_samples[i][j] + num_threads - 1) / num_threads;
             if (run_samples[j] < p_reps + rep_offset) {
               break_loop = false;
@@ -375,7 +376,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
           if (break_loop) {
             // Lock writing to this batch index in output_tensor.
             batch_locks[i].lock();
-            for (int j = 0; j < num_samples[i].size(); j++) {
+            for (size_t j = 0; j < num_samples[i].size(); j++) {
               rolling_sums[j] /= num_samples[i][j];
               (*output_tensor)(i, j) += static_cast<float>(rolling_sums[j]);
             }

tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc

@@ -40,6 +40,7 @@ limitations under the License.
 #include "tensorflow_quantum/core/ops/parse_context.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
 #include "tensorflow_quantum/core/src/circuit_parser_qsim.h"
+#include "tensorflow_quantum/core/src/util_balance_trajectory.h"
 #include "tensorflow_quantum/core/src/util_qsim.h"
 
 namespace tfq {
@@ -159,7 +160,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
     // 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 nescessary.
-    for (int i = 0; i < ncircuits.size(); i++) {
+    for (size_t i = 0; i < ncircuits.size(); i++) {
       int nq = num_qubits[i];
 
       if (nq > largest_nq) {
@@ -181,7 +182,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
         QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
                              sv, gathered_samples);
-        uint64_t q_ind = 0;
+        int q_ind = 0;
         uint64_t mask = 1;
         bool val = 0;
         while (q_ind < nq) {
@@ -252,7 +253,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
       auto local_gen = random_gen.ReserveSamples32(needed_random);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
 
-      for (int i = 0; i < ncircuits.size(); i++) {
+      for (size_t i = 0; i < ncircuits.size(); i++) {
         int nq = num_qubits[i];
         int j = start > 0 ? offset_prefix_sum[start - 1][i] : 0;
         int needed_samples = offset_prefix_sum[start][i] - j;
@@ -278,7 +279,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
           QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss,
                                sim, sv, gathered_samples);
 
-          uint64_t q_ind = 0;
+          int q_ind = 0;
           uint64_t mask = 1;
           bool val = 0;
           while (q_ind < nq) {