Change how this is done; just check input value

Author: mhucka

tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <stdlib.h>
 
+#include <limits>
 #include <string>
 
 #include "../qsim/lib/channel.h"
@@ -78,6 +79,12 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
     int num_samples = 0;
     OP_REQUIRES_OK(context, GetIndividualSample(context, &num_samples));
+    OP_REQUIRES(
+        context,
+        num_samples >= 0 && num_samples < std::numeric_limits<int>::max(),
+        tensorflow::errors::InvalidArgument(
+            absl::StrCat("num_samples must be between 0 and ",
+                         std::numeric_limits<int>::max(), ".")));
 
     // Construct qsim circuits.
     std::vector<NoisyQsimCircuit> qsim_circuits(programs.size(),
@@ -228,8 +235,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
         num_threads, std::vector<long>(output_dim_batch_size, 0));
 
     for (int i = 0; i < output_dim_batch_size; i++) {
-      uint64_t p_reps =
-          (static_cast<uint64_t>(num_samples) + num_threads - 1) / num_threads;
+      int p_reps = (num_samples + num_threads - 1) / num_threads;
       offset_prefix_sum[0][i] = rep_offsets[0][i] + p_reps;
       for (int j = 1; j < num_threads; j++) {
         offset_prefix_sum[j][i] += offset_prefix_sum[j - 1][i];

tensorflow_quantum/core/src/util_balance_trajectory.cc

@@ -75,10 +75,9 @@ void BalanceTrajectory(const int& num_samples, const int& num_threads,
 
   int prev_max_height = -1;
   for (size_t j = 0; j < (*thread_offsets)[0].size(); j++) {
-    uint64_t run_ceiling =
-        ((static_cast<uint64_t>(num_samples) + num_threads - 1) / num_threads);
-    uint64_t num_lo = num_threads * run_ceiling - num_samples;
-    uint64_t num_hi = num_threads - num_lo;
+    int run_ceiling = ((num_samples + num_threads - 1) / num_threads);
+    int num_lo = num_threads * run_ceiling - num_samples;
+    int num_hi = num_threads - num_lo;
     int cur_max = prev_max_height;
     for (int i = 0; i < num_threads; i++) {
       if (height[i] == cur_max && num_lo) {