From 2dfe2a87a7fa434e497579d306bff548aed4cf2e Mon Sep 17 00:00:00 2001
From: mhucka <mhucka@google.com>
Date: Mon, 22 Jun 2026 04:39:52 +0000
Subject: [PATCH 1/3] Fix vulnerability due to unchecked sample sizes

Fix for b/510431411
---
 .../ops/noise/noisy_expectation_op_test.py    | 10 ++++
 .../noisy_sampled_expectation_op_test.py      | 10 ++++
 .../core/ops/noise/noisy_samples_op_test.py   |  7 +++
 .../core/ops/noise/tfq_noisy_expectation.cc   | 35 ++++++-------
 .../noise/tfq_noisy_sampled_expectation.cc    | 34 ++++++-------
 .../core/ops/noise/tfq_noisy_samples.cc       | 51 +++++++++++--------
 tensorflow_quantum/core/ops/parse_context.cc  | 33 +++++++-----
 .../core/ops/tfq_simulate_ops_test.py         |  9 ++++
 .../core/src/util_balance_trajectory.cc       | 12 ++---
 .../core/src/util_balance_trajectory_test.cc  | 16 ++++++
 10 files changed, 143 insertions(+), 74 deletions(-)

diff --git a/tensorflow_quantum/core/ops/noise/noisy_expectation_op_test.py b/tensorflow_quantum/core/ops/noise/noisy_expectation_op_test.py
index 953829318..249d57325 100644
--- a/tensorflow_quantum/core/ops/noise/noisy_expectation_op_test.py
+++ b/tensorflow_quantum/core/ops/noise/noisy_expectation_op_test.py
@@ -213,6 +213,16 @@ def test_noisy_expectation_inputs(self):
                 [[-1]] * batch_size)
             # pylint: enable=too-many-function-args
 
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'less than or equal to 10,000,000'):
+            # pylint: disable=too-many-function-args
+            noisy_expectation_op.expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]),
+                [[20000000]] * batch_size)
+            # pylint: enable=too-many-function-args
+
         with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                     expected_regex='do not match'):
             # wrong symbol_values size.
diff --git a/tensorflow_quantum/core/ops/noise/noisy_sampled_expectation_op_test.py b/tensorflow_quantum/core/ops/noise/noisy_sampled_expectation_op_test.py
index c76122070..4d320d7b5 100644
--- a/tensorflow_quantum/core/ops/noise/noisy_sampled_expectation_op_test.py
+++ b/tensorflow_quantum/core/ops/noise/noisy_sampled_expectation_op_test.py
@@ -213,6 +213,16 @@ def test_noisy_expectation_inputs(self):
                 [[-1]] * batch_size)
             # pylint: enable=too-many-function-args
 
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'less than or equal to 10,000,000'):
+            # pylint: disable=too-many-function-args
+            noisy_sampled_expectation_op.sampled_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]),
+                [[20000000]] * batch_size)
+            # pylint: enable=too-many-function-args
+
         with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                     expected_regex='do not match'):
             # wrong symbol_values size.
diff --git a/tensorflow_quantum/core/ops/noise/noisy_samples_op_test.py b/tensorflow_quantum/core/ops/noise/noisy_samples_op_test.py
index 7790bda8b..b3c4707b0 100644
--- a/tensorflow_quantum/core/ops/noise/noisy_samples_op_test.py
+++ b/tensorflow_quantum/core/ops/noise/noisy_samples_op_test.py
@@ -133,6 +133,13 @@ def test_simulate_samples_inputs(self):
                                      symbol_names, symbol_values_array,
                                      ['junk'])
 
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'must be between 0 and 10,000,000'):
+            # num_samples is too large (above capacity limit).
+            noisy_samples_op.samples(util.convert_to_tensor(circuit_batch),
+                                     symbol_names, symbol_values_array,
+                                     [20000000])
+
         with self.assertRaisesRegex(TypeError, 'missing'):
             # too few tensors.
             # pylint: disable=no-value-for-parameter
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
index de32b8806..6d835f6b3 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
@@ -119,8 +119,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
     Status parse_status = ::tensorflow::Status();
     auto p_lock = absl::Mutex();
-    auto construct_f = [&](int64_t start, int64_t end) {
-      for (int64_t i = start; i < end; i++) {
+    auto construct_f = [&](int start, int end) {
+      for (int i = start; i < end; i++) {
         Status local = NoisyQsimCircuitFromProgram(
             programs[i], maps[i], num_qubits[i], false, &qsim_circuits[i]);
         NESTED_FN_STATUS_SYNC(parse_status, local, p_lock);
@@ -132,8 +132,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
         programs.size(), num_cycles, construct_f);
     OP_REQUIRES_OK(context, parse_status);
 
-    uint64_t max_num_qubits = 0;
-    for (const uint64_t num : num_qubits) {
+    int max_num_qubits = 0;
+    for (const int num : num_qubits) {
       max_num_qubits = std::max(max_num_qubits, num);
     }
 
@@ -173,7 +173,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
                                          qsim::MultiQubitGateFuser, Simulator>;
 
     // Begin simulation.
-    uint64_t largest_nq = 1;
+    int largest_nq = 1;
     Simulator sim = Simulator(tfq_for);
     StateSpace ss = StateSpace(tfq_for);
     auto sv = ss.Create(largest_nq);
@@ -187,15 +187,16 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
       }
     }
     random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
+    int64_t num_samples_needed = static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1;
     auto local_gen =
-        random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
+        random_gen.ReserveSamples128(num_samples_needed);
     tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // 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 (size_t i = 0; i < ncircuits.size(); i++) {
-      uint64_t nq = num_qubits[i];
+      int nq = num_qubits[i];
 
       // (#679) Just ignore empty program
       if (ncircuits[i].channels.size() == 0) {
@@ -256,7 +257,7 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
   }
 
   void ComputeSmall(const std::vector<int>& num_qubits,
-                    const uint64_t max_num_qubits,
+                    const int max_num_qubits,
                     const std::vector<NoisyQsimCircuit>& ncircuits,
                     const std::vector<std::vector<PauliSum>>& pauli_sums,
                     const std::vector<std::vector<int>>& num_samples,
@@ -294,23 +295,23 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
     Status compute_status = ::tensorflow::Status();
     auto c_lock = absl::Mutex();
-    auto DoWork = [&](int64_t start, int64_t end) {
+    auto DoWork = [&](int start, int end) {
       // Begin simulation.
       const auto tfq_for = qsim::SequentialFor(1);
-      uint64_t largest_nq = 1;
+      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);
 
-      int n_rand = ncircuits.size() * max_n_shots + 1;
+      int64_t n_rand = static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1;
       n_rand = (n_rand + num_threads) / num_threads;
       auto local_gen =
-          random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
+          random_gen.ReserveSamples128(static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (size_t i = 0; i < ncircuits.size(); i++) {
-        uint64_t nq = num_qubits[i];
+        int nq = num_qubits[i];
         int rep_offset = rep_offsets[start][i];
 
         // (#679) Just ignore empty program
@@ -343,8 +344,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
           // Compute expectations across all ops using this trajectory.
           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) {
+            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            if (static_cast<int64_t>(run_samples[j]) >= p_reps + rep_offset) {
               continue;
             }
             float exp_v = 0.0;
@@ -360,8 +361,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
           // Check if we have run enough trajectories for all ops.
           bool break_loop = true;
           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) {
+            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            if (static_cast<int64_t>(run_samples[j]) < p_reps + rep_offset) {
               break_loop = false;
               break;
             }
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
index 9b9058165..16c8d7ba9 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
@@ -120,8 +120,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
 
     Status parse_status = ::tensorflow::Status();
     auto p_lock = absl::Mutex();
-    auto construct_f = [&](int64_t start, int64_t end) {
-      for (int64_t i = start; i < end; i++) {
+    auto construct_f = [&](int start, int end) {
+      for (int i = start; i < end; i++) {
         Status local = NoisyQsimCircuitFromProgram(
             programs[i], maps[i], num_qubits[i], false, &qsim_circuits[i]);
         NESTED_FN_STATUS_SYNC(parse_status, local, p_lock);
@@ -133,8 +133,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
         programs.size(), num_cycles, construct_f);
     OP_REQUIRES_OK(context, parse_status);
 
-    uint64_t max_num_qubits = 0;
-    for (const uint64_t num : num_qubits) {
+    int max_num_qubits = 0;
+    for (const int num : num_qubits) {
       max_num_qubits = std::max(max_num_qubits, num);
     }
 
@@ -174,7 +174,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
                                          qsim::MultiQubitGateFuser, Simulator>;
 
     // Begin simulation.
-    uint64_t largest_nq = 1;
+    int largest_nq = 1;
     Simulator sim = Simulator(tfq_for);
     StateSpace ss = StateSpace(tfq_for);
     auto sv = ss.Create(largest_nq);
@@ -191,15 +191,15 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
       }
     }
     random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
-    auto local_gen = random_gen.ReserveSamples128(
-        ncircuits.size() * (1 + max_psum_length) * max_n_shots);
+    int64_t num_samples_needed = static_cast<int64_t>(ncircuits.size()) * (1 + max_psum_length) * max_n_shots;
+    auto local_gen = random_gen.ReserveSamples128(num_samples_needed);
     tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // 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 (size_t i = 0; i < ncircuits.size(); i++) {
-      uint64_t nq = num_qubits[i];
+      int nq = num_qubits[i];
 
       // (#679) Just ignore empty program
       if (ncircuits[i].channels.empty()) {
@@ -260,7 +260,7 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
   }
 
   void ComputeSmall(const std::vector<int>& num_qubits,
-                    const uint64_t max_num_qubits,
+                    const int max_num_qubits,
                     const std::vector<NoisyQsimCircuit>& ncircuits,
                     const std::vector<std::vector<PauliSum>>& pauli_sums,
                     const std::vector<std::vector<int>>& num_samples,
@@ -301,22 +301,22 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
 
     Status compute_status = ::tensorflow::Status();
     auto c_lock = absl::Mutex();
-    auto DoWork = [&](int64_t start, int64_t end) {
+    auto DoWork = [&](int start, int end) {
       // Begin simulation.
       const auto tfq_for = qsim::SequentialFor(1);
-      uint64_t largest_nq = 1;
+      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);
 
-      int num_rand = ncircuits.size() * (1 + max_psum_length) * max_n_shots;
+      int64_t num_rand = static_cast<int64_t>(ncircuits.size()) * (1 + max_psum_length) * max_n_shots;
       num_rand = (num_rand + num_threads) / num_threads + 1;
       auto local_gen = random_gen.ReserveSamples128(num_rand);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (size_t i = 0; i < ncircuits.size(); i++) {
-        uint64_t nq = num_qubits[i];
+        int nq = num_qubits[i];
         int rep_offset = rep_offsets[start][i];
 
         // (#679) Just ignore empty program
@@ -349,8 +349,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
 
           // Compute expectations across all ops using this trajectory.
           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) {
+            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            if (static_cast<int64_t>(run_samples[j]) >= p_reps + rep_offset) {
               continue;
             }
             float exp_v = 0.0;
@@ -366,8 +366,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
           // Check if we have run enough trajectories for all ops.
           bool break_loop = true;
           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) {
+            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            if (static_cast<int64_t>(run_samples[j]) < p_reps + rep_offset) {
               break_loop = false;
               break;
             }
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
index 6697f5071..927192bc8 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
+++ b/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"
@@ -60,7 +61,10 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
   void Compute(tensorflow::OpKernelContext* context) override {
     // TODO (mbbrough): add more dimension checks for other inputs here.
-    DCHECK_EQ(4, context->num_inputs());
+    const int num_inputs = context->num_inputs();
+    OP_REQUIRES(context, num_inputs == 4,
+                tensorflow::errors::InvalidArgument(absl::StrCat(
+                    "Expected 4 inputs, got ", num_inputs, " inputs.")));
 
     // Parse to Program Proto and num_qubits.
     std::vector<Program> programs;
@@ -79,6 +83,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 <= 10000000,
+        tensorflow::errors::InvalidArgument(
+            "num_samples must be between 0 and 10,000,000. Got: ",
+            num_samples));
 
     // Construct qsim circuits.
     std::vector<NoisyQsimCircuit> qsim_circuits(programs.size(),
@@ -86,8 +96,8 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
     Status parse_status = ::tensorflow::Status();
     auto p_lock = absl::Mutex();
-    auto construct_f = [&](int64_t start, int64_t end) {
-      for (int64_t i = start; i < end; i++) {
+    auto construct_f = [&](int start, int end) {
+      for (int i = start; i < end; i++) {
         auto r = NoisyQsimCircuitFromProgram(
             programs[i], maps[i], num_qubits[i], true, &qsim_circuits[i]);
         NESTED_FN_STATUS_SYNC(parse_status, r, p_lock);
@@ -99,12 +109,12 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
         programs.size(), num_cycles, construct_f);
     OP_REQUIRES_OK(context, parse_status);
 
-    uint64_t max_num_qubits = 0;
-    for (const uint64_t num : num_qubits) {
+    int max_num_qubits = 0;
+    for (const int num : num_qubits) {
       max_num_qubits = std::max(max_num_qubits, num);
     }
 
-    const size_t output_dim_size = maps.size();
+    const int output_dim_size = maps.size();
     tensorflow::TensorShape output_shape;
     output_shape.AddDim(output_dim_size);
     output_shape.AddDim(num_samples);
@@ -134,7 +144,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
  private:
   void ComputeLarge(const std::vector<int>& num_qubits,
-                    const uint64_t max_num_qubits, const int num_samples,
+                    const int max_num_qubits, const int num_samples,
                     const std::vector<NoisyQsimCircuit>& ncircuits,
                     tensorflow::OpKernelContext* context,
                     tensorflow::TTypes<int8_t, 3>::Tensor* output_tensor) {
@@ -147,22 +157,23 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
                                          qsim::MultiQubitGateFuser, Simulator>;
 
     // Begin simulation.
-    uint64_t largest_nq = 1;
+    int largest_nq = 1;
     Simulator sim = Simulator(tfq_for);
     StateSpace ss = StateSpace(tfq_for);
     auto sv = ss.Create(largest_nq);
 
     tensorflow::GuardedPhiloxRandom random_gen;
     random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
+    int64_t num_samples_needed = 2 * static_cast<int64_t>(num_samples) * ncircuits.size() + 2;
     auto local_gen =
-        random_gen.ReserveSamples32(2 * num_samples * ncircuits.size() + 2);
+        random_gen.ReserveSamples32(num_samples_needed);
     tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // 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.
+    // a larger circuit we will grow the Statevector as nescessary.
     for (size_t i = 0; i < ncircuits.size(); i++) {
-      uint64_t nq = num_qubits[i];
+      int nq = num_qubits[i];
 
       if (nq > largest_nq) {
         // need to switch to larger statespace.
@@ -183,7 +194,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) {
@@ -203,7 +214,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
   }
 
   void ComputeSmall(const std::vector<int>& num_qubits,
-                    const uint64_t max_num_qubits, const int num_samples,
+                    const int max_num_qubits, const int num_samples,
                     const std::vector<NoisyQsimCircuit>& ncircuits,
                     tensorflow::OpKernelContext* context,
                     tensorflow::TTypes<int8_t, 3>::Tensor* output_tensor) {
@@ -229,7 +240,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++) {
-      int p_reps = (num_samples + num_threads - 1) / num_threads;
+      int64_t p_reps = (static_cast<int64_t>(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];
@@ -240,22 +251,22 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
     tensorflow::GuardedPhiloxRandom random_gen;
     random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
 
-    auto DoWork = [&](int64_t start, int64_t end) {
+    auto DoWork = [&](int start, int end) {
       // Begin simulation.
       const auto tfq_for = qsim::SequentialFor(1);
-      uint64_t largest_nq = 1;
+      int largest_nq = 1;
       Simulator sim = Simulator(tfq_for);
       StateSpace ss = StateSpace(tfq_for);
       auto sv = ss.Create(largest_nq);
 
-      int needed_random =
-          4 * (num_samples * ncircuits.size() + num_threads) / num_threads;
+      int64_t needed_random =
+          4 * (static_cast<int64_t>(num_samples) * ncircuits.size() + num_threads) / num_threads;
       needed_random += 4;
       auto local_gen = random_gen.ReserveSamples32(needed_random);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (size_t i = 0; i < ncircuits.size(); i++) {
-        uint64_t nq = num_qubits[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;
         if (needed_samples <= 0) {
@@ -280,7 +291,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) {
diff --git a/tensorflow_quantum/core/ops/parse_context.cc b/tensorflow_quantum/core/ops/parse_context.cc
index 04f6324d5..855ce3e9a 100644
--- a/tensorflow_quantum/core/ops/parse_context.cc
+++ b/tensorflow_quantum/core/ops/parse_context.cc
@@ -77,8 +77,8 @@ Status ParsePrograms(OpKernelContext* context, const std::string& input_name,
   const int num_programs = program_strings.dimension(0);
   programs->assign(num_programs, Program());
 
-  auto DoWork = [&](int64_t start, int64_t end) {
-    for (int64_t i = start; i < end; i++) {
+  auto DoWork = [&](int start, int end) {
+    for (int i = start; i < end; i++) {
       OP_REQUIRES_OK(context, ParseProto(program_strings(i), &programs->at(i)));
     }
   };
@@ -112,8 +112,8 @@ Status ParsePrograms2D(OpKernelContext* context, const std::string& input_name,
   const int num_entries = program_strings.dimension(1);
   programs->assign(num_programs, std::vector<Program>(num_entries, Program()));
 
-  auto DoWork = [&](int64_t start, int64_t end) {
-    for (int64_t i = start; i < end; i++) {
+  auto DoWork = [&](int start, int end) {
+    for (int i = start; i < end; i++) {
       OP_REQUIRES_OK(
           context,
           ParseProto(program_strings(i / num_entries, i % num_entries),
@@ -181,8 +181,8 @@ Status GetProgramsAndNumQubits(
 
   // Resolve qubit ID's in parallel.
   num_qubits->assign(programs->size(), -1);
-  auto DoWork = [&](int64_t start, int64_t end) {
-    for (int64_t i = start; i < end; i++) {
+  auto DoWork = [&](int start, int end) {
+    for (int i = start; i < end; i++) {
       Program& program = (*programs)[i];
       unsigned int this_num_qubits;
       if (p_sums) {
@@ -232,8 +232,8 @@ tensorflow::Status GetProgramsAndNumQubits(
 
   // Resolve qubit ID's in parallel.
   num_qubits->assign(programs->size(), -1);
-  auto DoWork = [&](int64_t start, int64_t end) {
-    for (int64_t i = start; i < end; i++) {
+  auto DoWork = [&](int start, int end) {
+    for (int i = start; i < end; i++) {
       Program& program = (*programs)[i];
       unsigned int this_num_qubits;
       OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits,
@@ -270,8 +270,8 @@ Status GetPauliSums(OpKernelContext* context,
   p_sums->assign(sum_specs.dimension(0),
                  std::vector<PauliSum>(sum_specs.dimension(1), PauliSum()));
   const int op_dim = sum_specs.dimension(1);
-  auto DoWork = [&](int64_t start, int64_t end) {
-    for (int64_t ii = start; ii < end; ii++) {
+  auto DoWork = [&](int start, int end) {
+    for (int ii = start; ii < end; ii++) {
       const int i = ii / op_dim;
       const int j = ii % op_dim;
       PauliSum p;
@@ -328,8 +328,8 @@ Status GetSymbolMaps(OpKernelContext* context, std::vector<SymbolMap>* maps) {
   maps->assign(symbol_values.dimension(0), SymbolMap());
 
   const int symbol_dim = symbol_values.dimension(1);
-  auto DoWork = [&](int64_t start, int64_t end) {
-    for (int64_t i = start; i < end; i++) {
+  auto DoWork = [&](int start, int end) {
+    for (int i = start; i < end; i++) {
       for (int j = 0; j < symbol_dim; j++) {
         const std::string& name = symbol_names(j);
         const float value = symbol_values(i, j);
@@ -369,10 +369,10 @@ tensorflow::Status GetNumSamples(
     sub_parsed_num_samples.reserve(matrix_num_samples.dimension(1));
     for (unsigned int j = 0; j < matrix_num_samples.dimension(1); j++) {
       const int num_samples = matrix_num_samples(i, j);
-      if (num_samples < 1) {
+      if (num_samples < 1 || num_samples > 10000000) {
         return Status(static_cast<tensorflow::errors::Code>(
                           absl::StatusCode::kInvalidArgument),
-                      "Each element of num_samples must be greater than 0.");
+                      "Each element of num_samples must be greater than 0 and less than or equal to 10,000,000.");
       }
       sub_parsed_num_samples.push_back(num_samples);
     }
@@ -408,6 +408,11 @@ Status GetIndividualSample(tensorflow::OpKernelContext* context,
   }
 
   (*n_samples) = vector_num_samples(0);
+  if ((*n_samples) < 0 || (*n_samples) > 10000000) {
+    return Status(static_cast<tensorflow::errors::Code>(
+                      absl::StatusCode::kInvalidArgument),
+                  "num_samples must be between 0 and 10,000,000.");
+  }
   return ::tensorflow::Status();
 }
 
diff --git a/tensorflow_quantum/core/ops/tfq_simulate_ops_test.py b/tensorflow_quantum/core/ops/tfq_simulate_ops_test.py
index 93c1770e2..3173e6c68 100644
--- a/tensorflow_quantum/core/ops/tfq_simulate_ops_test.py
+++ b/tensorflow_quantum/core/ops/tfq_simulate_ops_test.py
@@ -672,6 +672,15 @@ def test_simulate_sampled_expectation_inputs(self):
                 util.convert_to_tensor([[x] for x in pauli_sums]),
                 [[-1]] * batch_size)
 
+        with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                    'less than or equal to 10,000,000'):
+            # pylint: disable=too-many-function-args
+            tfq_simulate_ops.tfq_simulate_sampled_expectation(
+                util.convert_to_tensor(circuit_batch), symbol_names,
+                symbol_values_array,
+                util.convert_to_tensor([[x] for x in pauli_sums]),
+                [[20000000]] * batch_size)
+
         with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
                                     expected_regex='do not match'):
             # wrong symbol_values size.
diff --git a/tensorflow_quantum/core/src/util_balance_trajectory.cc b/tensorflow_quantum/core/src/util_balance_trajectory.cc
index 8351e49b1..3e2dac8f2 100644
--- a/tensorflow_quantum/core/src/util_balance_trajectory.cc
+++ b/tensorflow_quantum/core/src/util_balance_trajectory.cc
@@ -36,9 +36,9 @@ void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
   }
   int prev_max_height = -1;
   for (size_t j = 0; j < num_samples.size(); j++) {
-    int run_ceiling = ((rep_limits[j] + num_threads - 1) / num_threads);
-    int num_lo = num_threads * run_ceiling - rep_limits[j];
-    int num_hi = num_threads - num_lo;
+    int64_t run_ceiling = ((static_cast<int64_t>(rep_limits[j]) + num_threads - 1) / num_threads);
+    int64_t num_lo = num_threads * run_ceiling - rep_limits[j];
+    int64_t 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) {
@@ -75,9 +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++) {
-    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;
+    int64_t run_ceiling = ((static_cast<int64_t>(num_samples) + num_threads - 1) / num_threads);
+    int64_t num_lo = num_threads * run_ceiling - num_samples;
+    int64_t 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) {
diff --git a/tensorflow_quantum/core/src/util_balance_trajectory_test.cc b/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
index f361f5754..46a06734f 100644
--- a/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
+++ b/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
@@ -140,5 +140,21 @@ TEST(UtilQsimTest, BalanceTrajectory1D_2) {
   AssertWellBalanced(tmp, num_threads, offsets);
 }
 
+TEST(UtilQsimTest, BalanceTrajectoryOverflowPrevention) {
+  const int n_reps = 2147483640;
+  const int num_threads = 4;
+  std::vector<std::vector<int>> offsets = {
+      {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}, {0, 0, 0, 0, 0}};
+
+  BalanceTrajectory(n_reps, num_threads, &offsets);
+
+  // Since num_lo is 0 and num_hi is 4, all offsets should remain 0 and not overflow/become negative
+  for (int i = 0; i < num_threads; i++) {
+    for (size_t j = 0; j < offsets[i].size(); j++) {
+      EXPECT_EQ(offsets[i][j], 0);
+    }
+  }
+}
+
 }  // namespace
 }  // namespace tfq

From 42188d27c02edcadf7a3c5d9cf1bbebbc471b2cb Mon Sep 17 00:00:00 2001
From: mhucka <mhucka@google.com>
Date: Mon, 22 Jun 2026 04:41:11 +0000
Subject: [PATCH 2/3] Fix typo

---
 tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
index 927192bc8..9d42cf2bb 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
@@ -171,7 +171,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.
+    // a larger circuit we will grow the Statevector as necessary.
     for (size_t i = 0; i < ncircuits.size(); i++) {
       int nq = num_qubits[i];
 

From 4e354dfa362bd75bff54c06562bbb7dcf00924a8 Mon Sep 17 00:00:00 2001
From: mhucka <mhucka@google.com>
Date: Mon, 22 Jun 2026 04:47:07 +0000
Subject: [PATCH 3/3] Format

---
 .../core/ops/noise/tfq_noisy_expectation.cc   | 18 +++++++++------
 .../noise/tfq_noisy_sampled_expectation.cc    | 14 +++++++----
 .../core/ops/noise/tfq_noisy_samples.cc       | 23 ++++++++++---------
 tensorflow_quantum/core/ops/parse_context.cc  |  3 ++-
 .../core/src/util_balance_trajectory.cc       |  6 +++--
 .../core/src/util_balance_trajectory_test.cc  |  3 ++-
 6 files changed, 41 insertions(+), 26 deletions(-)

diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
index 6d835f6b3..c3d4abada 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc
@@ -187,9 +187,9 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
       }
     }
     random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
-    int64_t num_samples_needed = static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1;
-    auto local_gen =
-        random_gen.ReserveSamples128(num_samples_needed);
+    int64_t num_samples_needed =
+        static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1;
+    auto local_gen = random_gen.ReserveSamples128(num_samples_needed);
     tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // Simulate programs one by one. Parallelizing over state vectors
@@ -306,8 +306,8 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
       int64_t n_rand = static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1;
       n_rand = (n_rand + num_threads) / num_threads;
-      auto local_gen =
-          random_gen.ReserveSamples128(static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1);
+      auto local_gen = random_gen.ReserveSamples128(
+          static_cast<int64_t>(ncircuits.size()) * max_n_shots + 1);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
 
       for (size_t i = 0; i < ncircuits.size(); i++) {
@@ -344,7 +344,9 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
 
           // Compute expectations across all ops using this trajectory.
           for (size_t j = 0; j < pauli_sums[i].size(); j++) {
-            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            int64_t p_reps =
+                (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) /
+                num_threads;
             if (static_cast<int64_t>(run_samples[j]) >= p_reps + rep_offset) {
               continue;
             }
@@ -361,7 +363,9 @@ class TfqNoisyExpectationOp : public tensorflow::OpKernel {
           // Check if we have run enough trajectories for all ops.
           bool break_loop = true;
           for (size_t j = 0; j < num_samples[i].size(); j++) {
-            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            int64_t p_reps =
+                (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) /
+                num_threads;
             if (static_cast<int64_t>(run_samples[j]) < p_reps + rep_offset) {
               break_loop = false;
               break;
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
index 16c8d7ba9..e6f22fa85 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_sampled_expectation.cc
@@ -191,7 +191,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
       }
     }
     random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
-    int64_t num_samples_needed = static_cast<int64_t>(ncircuits.size()) * (1 + max_psum_length) * max_n_shots;
+    int64_t num_samples_needed = static_cast<int64_t>(ncircuits.size()) *
+                                 (1 + max_psum_length) * max_n_shots;
     auto local_gen = random_gen.ReserveSamples128(num_samples_needed);
     tensorflow::random::SimplePhilox rand_source(&local_gen);
 
@@ -310,7 +311,8 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
       auto sv = ss.Create(largest_nq);
       auto scratch = ss.Create(largest_nq);
 
-      int64_t num_rand = static_cast<int64_t>(ncircuits.size()) * (1 + max_psum_length) * max_n_shots;
+      int64_t num_rand = static_cast<int64_t>(ncircuits.size()) *
+                         (1 + max_psum_length) * max_n_shots;
       num_rand = (num_rand + num_threads) / num_threads + 1;
       auto local_gen = random_gen.ReserveSamples128(num_rand);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
@@ -349,7 +351,9 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
 
           // Compute expectations across all ops using this trajectory.
           for (size_t j = 0; j < pauli_sums[i].size(); j++) {
-            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            int64_t p_reps =
+                (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) /
+                num_threads;
             if (static_cast<int64_t>(run_samples[j]) >= p_reps + rep_offset) {
               continue;
             }
@@ -366,7 +370,9 @@ class TfqNoisySampledExpectationOp : public tensorflow::OpKernel {
           // Check if we have run enough trajectories for all ops.
           bool break_loop = true;
           for (size_t j = 0; j < num_samples[i].size(); j++) {
-            int64_t p_reps = (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) / num_threads;
+            int64_t p_reps =
+                (static_cast<int64_t>(num_samples[i][j]) + num_threads - 1) /
+                num_threads;
             if (static_cast<int64_t>(run_samples[j]) < p_reps + rep_offset) {
               break_loop = false;
               break;
diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
index 9d42cf2bb..8519d62e6 100644
--- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
+++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
@@ -83,12 +83,10 @@ 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 <= 10000000,
-        tensorflow::errors::InvalidArgument(
-            "num_samples must be between 0 and 10,000,000. Got: ",
-            num_samples));
+    OP_REQUIRES(context, num_samples >= 0 && num_samples <= 10000000,
+                tensorflow::errors::InvalidArgument(
+                    "num_samples must be between 0 and 10,000,000. Got: ",
+                    num_samples));
 
     // Construct qsim circuits.
     std::vector<NoisyQsimCircuit> qsim_circuits(programs.size(),
@@ -164,9 +162,9 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
 
     tensorflow::GuardedPhiloxRandom random_gen;
     random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
-    int64_t num_samples_needed = 2 * static_cast<int64_t>(num_samples) * ncircuits.size() + 2;
-    auto local_gen =
-        random_gen.ReserveSamples32(num_samples_needed);
+    int64_t num_samples_needed =
+        2 * static_cast<int64_t>(num_samples) * ncircuits.size() + 2;
+    auto local_gen = random_gen.ReserveSamples32(num_samples_needed);
     tensorflow::random::SimplePhilox rand_source(&local_gen);
 
     // Simulate programs one by one. Parallelizing over state vectors
@@ -240,7 +238,8 @@ 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++) {
-      int64_t p_reps = (static_cast<int64_t>(num_samples) + num_threads - 1) / num_threads;
+      int64_t p_reps =
+          (static_cast<int64_t>(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];
@@ -260,7 +259,9 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {
       auto sv = ss.Create(largest_nq);
 
       int64_t needed_random =
-          4 * (static_cast<int64_t>(num_samples) * ncircuits.size() + num_threads) / num_threads;
+          4 *
+          (static_cast<int64_t>(num_samples) * ncircuits.size() + num_threads) /
+          num_threads;
       needed_random += 4;
       auto local_gen = random_gen.ReserveSamples32(needed_random);
       tensorflow::random::SimplePhilox rand_source(&local_gen);
diff --git a/tensorflow_quantum/core/ops/parse_context.cc b/tensorflow_quantum/core/ops/parse_context.cc
index 855ce3e9a..a3ccf28ae 100644
--- a/tensorflow_quantum/core/ops/parse_context.cc
+++ b/tensorflow_quantum/core/ops/parse_context.cc
@@ -372,7 +372,8 @@ tensorflow::Status GetNumSamples(
       if (num_samples < 1 || num_samples > 10000000) {
         return Status(static_cast<tensorflow::errors::Code>(
                           absl::StatusCode::kInvalidArgument),
-                      "Each element of num_samples must be greater than 0 and less than or equal to 10,000,000.");
+                      "Each element of num_samples must be greater than 0 and "
+                      "less than or equal to 10,000,000.");
       }
       sub_parsed_num_samples.push_back(num_samples);
     }
diff --git a/tensorflow_quantum/core/src/util_balance_trajectory.cc b/tensorflow_quantum/core/src/util_balance_trajectory.cc
index 3e2dac8f2..1f5adb30c 100644
--- a/tensorflow_quantum/core/src/util_balance_trajectory.cc
+++ b/tensorflow_quantum/core/src/util_balance_trajectory.cc
@@ -36,7 +36,8 @@ void BalanceTrajectory(const std::vector<std::vector<int>>& num_samples,
   }
   int prev_max_height = -1;
   for (size_t j = 0; j < num_samples.size(); j++) {
-    int64_t run_ceiling = ((static_cast<int64_t>(rep_limits[j]) + num_threads - 1) / num_threads);
+    int64_t run_ceiling =
+        ((static_cast<int64_t>(rep_limits[j]) + num_threads - 1) / num_threads);
     int64_t num_lo = num_threads * run_ceiling - rep_limits[j];
     int64_t num_hi = num_threads - num_lo;
     int cur_max = prev_max_height;
@@ -75,7 +76,8 @@ 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++) {
-    int64_t run_ceiling = ((static_cast<int64_t>(num_samples) + num_threads - 1) / num_threads);
+    int64_t run_ceiling =
+        ((static_cast<int64_t>(num_samples) + num_threads - 1) / num_threads);
     int64_t num_lo = num_threads * run_ceiling - num_samples;
     int64_t num_hi = num_threads - num_lo;
     int cur_max = prev_max_height;
diff --git a/tensorflow_quantum/core/src/util_balance_trajectory_test.cc b/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
index 46a06734f..6964d7668 100644
--- a/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
+++ b/tensorflow_quantum/core/src/util_balance_trajectory_test.cc
@@ -148,7 +148,8 @@ TEST(UtilQsimTest, BalanceTrajectoryOverflowPrevention) {
 
   BalanceTrajectory(n_reps, num_threads, &offsets);
 
-  // Since num_lo is 0 and num_hi is 4, all offsets should remain 0 and not overflow/become negative
+  // Since num_lo is 0 and num_hi is 4, all offsets should remain 0 and not
+  // overflow/become negative
   for (int i = 0; i < num_threads; i++) {
     for (size_t j = 0; j < offsets[i].size(); j++) {
       EXPECT_EQ(offsets[i][j], 0);