From ad801580095843f0a968fa5ad6e8c1bc9609650e Mon Sep 17 00:00:00 2001 From: mhucka Date: Fri, 22 May 2026 17:58:15 +0000 Subject: [PATCH 1/4] Use size_t or int64_t for trajectory size arithmetic --- tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc | 2 +- tensorflow_quantum/core/src/util_balance_trajectory.cc | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc index a09f826b9..04dc4e66b 100644 --- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc +++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc @@ -228,7 +228,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel { num_threads, std::vector(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; + uint64_t p_reps = (static_cast(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]; diff --git a/tensorflow_quantum/core/src/util_balance_trajectory.cc b/tensorflow_quantum/core/src/util_balance_trajectory.cc index 8351e49b1..8a81f45fe 100644 --- a/tensorflow_quantum/core/src/util_balance_trajectory.cc +++ b/tensorflow_quantum/core/src/util_balance_trajectory.cc @@ -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; + uint64_t run_ceiling = ((static_cast(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 cur_max = prev_max_height; for (int i = 0; i < num_threads; i++) { if (height[i] == cur_max && num_lo) { From f598888aac16abbc1d6e9d8c20a465620c9d2558 Mon Sep 17 00:00:00 2001 From: mhucka Date: Fri, 22 May 2026 18:09:25 +0000 Subject: [PATCH 2/4] Format --- tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc | 3 ++- tensorflow_quantum/core/src/util_balance_trajectory.cc | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc index 04dc4e66b..3f2c3f619 100644 --- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc +++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc @@ -228,7 +228,8 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel { num_threads, std::vector(output_dim_batch_size, 0)); for (int i = 0; i < output_dim_batch_size; i++) { - uint64_t p_reps = (static_cast(num_samples) + num_threads - 1) / num_threads; + uint64_t p_reps = + (static_cast(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]; diff --git a/tensorflow_quantum/core/src/util_balance_trajectory.cc b/tensorflow_quantum/core/src/util_balance_trajectory.cc index 8a81f45fe..189012793 100644 --- a/tensorflow_quantum/core/src/util_balance_trajectory.cc +++ b/tensorflow_quantum/core/src/util_balance_trajectory.cc @@ -75,7 +75,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++) { - uint64_t run_ceiling = ((static_cast(num_samples) + num_threads - 1) / num_threads); + uint64_t run_ceiling = + ((static_cast(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 cur_max = prev_max_height; From 991b01a677ab9383dcd120e87c8912edc56153f9 Mon Sep 17 00:00:00 2001 From: mhucka Date: Fri, 22 May 2026 20:27:02 +0000 Subject: [PATCH 3/4] Change how this is done; just check input value --- tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc | 10 ++++++++-- tensorflow_quantum/core/src/util_balance_trajectory.cc | 7 +++---- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc index 3f2c3f619..8b7f7d114 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 +#include #include #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::max(), + tensorflow::errors::InvalidArgument( + absl::StrCat("num_samples must be between 0 and ", + std::numeric_limits::max(), "."))); // Construct qsim circuits. std::vector qsim_circuits(programs.size(), @@ -228,8 +235,7 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel { num_threads, std::vector(output_dim_batch_size, 0)); for (int i = 0; i < output_dim_batch_size; i++) { - uint64_t p_reps = - (static_cast(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]; diff --git a/tensorflow_quantum/core/src/util_balance_trajectory.cc b/tensorflow_quantum/core/src/util_balance_trajectory.cc index 189012793..8351e49b1 100644 --- a/tensorflow_quantum/core/src/util_balance_trajectory.cc +++ b/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(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) { From ec3b0246791f123ebf74a6006e27d6b92f6c36f9 Mon Sep 17 00:00:00 2001 From: mhucka Date: Fri, 22 May 2026 20:33:34 +0000 Subject: [PATCH 4/4] Check the number of inputs to `Compute()` --- tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc | 5 ++++- 1 file changed, 4 insertions(+), 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 8b7f7d114..0e22956f0 100644 --- a/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc +++ b/tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc @@ -60,7 +60,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 programs;