Skip to content

Reject out-of-range num_samples in the op input parsers #1119

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mhucka wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Reject out-of-range num_samples in the op input parsers #1119

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f0fa7ed
Reject out-of-range num_samples in the op input parsers
mhucka Oct 13, 2026
c1d871a
Merge branch 'master' into fix-num-samples-bounds
mhucka Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions tensorflow_quantum/core/ops/math_ops/simulate_mps_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -437,6 +437,19 @@ def test_simulate_mps1d_samples_inputs(self):
symbol_names, symbol_values_array,
[[num_samples]])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'between 0 and .*Got 2147483647'):
# num_samples is above the supported maximum.
simulate_mps.mps_1d_sample(util.convert_to_tensor(circuit_batch),
symbol_names, symbol_values_array,
[2**31 - 1])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'between 0 and .*Got -1'):
# num_samples is negative.
simulate_mps.mps_1d_sample(util.convert_to_tensor(circuit_batch),
symbol_names, symbol_values_array, [-1])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'Could not find symbol in parameter map'):
# symbol_names tensor has the right type, but invalid value.
Expand Down Expand Up @@ -724,6 +737,15 @@ def test_simulate_mps1d_sampled_expectation_inputs(self):
util.convert_to_tensor([[x] for x in pauli_sums]),
num_samples[0])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'less than or equal to .*Got 2147483647'):
# num_samples is above the supported maximum.
simulate_mps.mps_1d_sampled_expectation(
util.convert_to_tensor(circuit_batch), symbol_names,
symbol_values_array,
util.convert_to_tensor([[x] for x in pauli_sums]),
[[2**31 - 1]] * batch_size)

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'Unparseable proto'):
# circuit tensor has the right type but invalid values.
Expand Down
9 changes: 9 additions & 0 deletions tensorflow_quantum/core/ops/noise/noisy_expectation_op_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,15 @@ 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 .*Got 2147483647'):
# num_samples is above the supported maximum.
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]),
[[2**31 - 1]] * batch_size)

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
expected_regex='do not match'):
# wrong symbol_values size.
Expand Down
9 changes: 9 additions & 0 deletions tensorflow_quantum/core/ops/noise/noisy_sampled_expectation_op_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,15 @@ 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 .*Got 2147483647'):
# num_samples is above the supported maximum.
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]),
[[2**31 - 1]] * batch_size)

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
expected_regex='do not match'):
# wrong symbol_values size.
Expand Down
28 changes: 28 additions & 0 deletions tensorflow_quantum/core/ops/noise/noisy_samples_op_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,34 @@ def test_simulate_samples_inputs(self):
symbol_names, symbol_values_array,
[[num_samples]])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'between 0 and .*Got 2147483647'):
# num_samples is above the supported maximum.
noisy_samples_op.samples(util.convert_to_tensor(circuit_batch),
symbol_names, symbol_values_array,
[2**31 - 1])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'between 0 and .*Got -1'):
# num_samples is negative.
noisy_samples_op.samples(util.convert_to_tensor(circuit_batch),
symbol_names, symbol_values_array, [-1])

# num_samples of 0 is still accepted and yields empty samples.
zero_samples = noisy_samples_op.samples(
util.convert_to_tensor(circuit_batch), symbol_names,
symbol_values_array, [0])
self.assertEqual(zero_samples.to_list(), [[]] * batch_size)

# The maximum itself (kMaxNumSamples) is accepted. An empty batch keeps
# the output empty, so this does not actually draw any samples.
empty_circuits = tf.raw_ops.Empty(shape=(0,), dtype=tf.string)
empty_symbols = tf.raw_ops.Empty(shape=(0,), dtype=tf.string)
empty_values = tf.raw_ops.Empty(shape=(0, 0), dtype=tf.float32)
max_samples = noisy_samples_op.samples(empty_circuits, empty_symbols,
empty_values, [10000000])
self.assertEqual(max_samples.shape[0], 0)

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'Unparseable proto'):
# programs tensor has the right type, but invalid value.
Expand Down
22 changes: 17 additions & 5 deletions tensorflow_quantum/core/ops/parse_context.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,10 +369,13 @@ 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) {
return Status(static_cast<tensorflow::errors::Code>(
absl::StatusCode::kInvalidArgument),
"Each element of num_samples must be greater than 0.");
if (num_samples < 1 || num_samples > kMaxNumSamples) {
return Status(
static_cast<tensorflow::errors::Code>(
absl::StatusCode::kInvalidArgument),
absl::StrCat("Each element of num_samples must be greater than 0 "
"and less than or equal to ",
kMaxNumSamples, ". Got ", num_samples, "."));
}
sub_parsed_num_samples.push_back(num_samples);
}
Expand Down Expand Up @@ -407,7 +410,16 @@ Status GetIndividualSample(tensorflow::OpKernelContext* context,
vector_num_samples.dimension(0), "."));
}

(*n_samples) = vector_num_samples(0);
const int num_samples = vector_num_samples(0);
if (num_samples < 0 || num_samples > kMaxNumSamples) {
return Status(
static_cast<tensorflow::errors::Code>(
absl::StatusCode::kInvalidArgument),
absl::StrCat("num_samples must be between 0 and ", kMaxNumSamples,
" inclusive. Got ", num_samples, "."));
}

(*n_samples) = num_samples;
return ::tensorflow::Status();
}

Expand Down
17 changes: 15 additions & 2 deletions tensorflow_quantum/core/ops/parse_context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,13 +96,26 @@ tensorflow::Status GetPauliSums(
tensorflow::Status GetSymbolMaps(tensorflow::OpKernelContext* context,
std::vector<SymbolMap>* maps);

// Parses the number of samples from the 'num_samples' input tensor.
// Upper bound on the number of samples accepted for any single circuit or
// pauli sum. The per-thread work split in BalanceTrajectory and in the noisy
// ops' ComputeSmall uses int arithmetic such as
// (num_samples + num_threads - 1) / num_threads, which overflows when
// num_samples is near INT_MAX; the resulting negative chunk sizes make the
// worker threads skip sampling and leave the output tensor uninitialized.
// 1e7 is far above practical shot counts while leaving more than 200x
// headroom to INT_MAX for any thread count.
inline constexpr int kMaxNumSamples = 10000000;

// Parses the number of samples from the 'num_samples' input tensor. Every
// element must be in the range [1, kMaxNumSamples].
tensorflow::Status GetNumSamples(
tensorflow::OpKernelContext* context,
std::vector<std::vector<int>>* parsed_num_samples);

// Parses the 'num_samples' input tensor when it is expected to only
// contain one element.
// contain one element. The element must be in the range
// [0, kMaxNumSamples]; 0 is legal because the sampling ops return an empty
// result for it.
tensorflow::Status GetIndividualSample(tensorflow::OpKernelContext* context,
int* n_samples);

Expand Down
32 changes: 32 additions & 0 deletions tensorflow_quantum/core/ops/tfq_simulate_ops_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -400,6 +400,20 @@ def test_simulate_samples_inputs(self):
util.convert_to_tensor(circuit_batch), symbol_names,
symbol_values_array, [[num_samples]])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'between 0 and .*Got 2147483647'):
# num_samples is above the supported maximum.
tfq_simulate_ops.tfq_simulate_samples(
util.convert_to_tensor(circuit_batch), symbol_names,
symbol_values_array, [2**31 - 1])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'between 0 and .*Got -1'):
# num_samples is negative.
tfq_simulate_ops.tfq_simulate_samples(
util.convert_to_tensor(circuit_batch), symbol_names,
symbol_values_array, [-1])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'Unparseable proto'):
# programs tensor has the right type, but invalid value.
Expand Down Expand Up @@ -672,6 +686,24 @@ 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 .*Got 2147483647'):
# num_samples is above the supported maximum.
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]),
[[2**31 - 1]] * batch_size)

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
'less than or equal to .*Got 2147483647'):
# Only the last element of num_samples is out of range.
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]),
[[10]] * (batch_size - 1) + [[2**31 - 1]])

with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
expected_regex='do not match'):
# wrong symbol_values size.
Expand Down
Loading