From 5aed85261ef9b2ee0f20b2740e802a6700f3ef50 Mon Sep 17 00:00:00 2001 From: Saksham Singh Rathore Date: Fri, 30 Jan 2026 20:48:48 +0530 Subject: [PATCH 1/4] [AutoGraph] Add warning when Python random module is used inside tf.function When Python's random module functions are used inside tf.function, the values are computed during tracing and become constants. This causes issues with XLA compilation when input shapes don't match the traced constant values. This change adds detection for Python random module functions and issues a helpful warning guiding users to use tf.random functions instead. Fixes #109111 --- .../python/autograph/converters/call_trees.py | 45 +++++++++++++++++ .../autograph/converters/call_trees_test.py | 48 +++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/tensorflow/python/autograph/converters/call_trees.py b/tensorflow/python/autograph/converters/call_trees.py index 3d694d45a17b0b..7ebdc4b40b8ba9 100644 --- a/tensorflow/python/autograph/converters/call_trees.py +++ b/tensorflow/python/autograph/converters/call_trees.py @@ -32,6 +32,31 @@ # TODO(mdan): Rename to FunctionCallsTransformer. +# Python random module functions that produce values traced as constants. +# These cause issues when used inside tf.function because the random value +# is computed during tracing and becomes a static constant in the graph. +_PYTHON_RANDOM_FUNCTIONS = frozenset([ + 'random.random', + 'random.randint', + 'random.randrange', + 'random.uniform', + 'random.choice', + 'random.choices', + 'random.sample', + 'random.shuffle', + 'random.gauss', + 'random.normalvariate', + 'random.betavariate', + 'random.expovariate', + 'random.gammavariate', + 'random.lognormvariate', + 'random.vonmisesvariate', + 'random.paretovariate', + 'random.weibullvariate', + 'random.triangular', + 'random.getrandbits', +]) + class _Function(object): @@ -42,6 +67,7 @@ def __init__(self): set_trace_warned = False +python_random_warned = False class _ArgTemplateBuilder(object): @@ -187,6 +213,25 @@ def visit_Call(self, node): set_trace_warned = True return node + # Warn when Python random module functions are used inside tf.function. + # These values are computed at trace time and become constants in the graph, + # which can cause shape mismatches or unexpected behavior at runtime, + # especially with XLA compilation. + if full_name in _PYTHON_RANDOM_FUNCTIONS: + global python_random_warned + if not python_random_warned: + ag_logging.warning( + 'Detected use of Python\'s `%s()` inside a tf.function. ' + 'The random value is computed during tracing and becomes a ' + 'constant in the graph, which may cause shape mismatches or ' + 'unexpected behavior, especially with XLA compilation. ' + 'Use `tf.random` functions instead for dynamic random values. ' + 'For example, replace `random.randint(a, b)` with ' + '`tf.random.uniform([], a, b, dtype=tf.int32)`. ' + 'See https://www.tensorflow.org/guide/function#executing_python_side_effects', + full_name) + python_random_warned = True + if (full_name == 'print' and not self.ctx.user.options.uses(converter.Feature.BUILTIN_FUNCTIONS)): return node diff --git a/tensorflow/python/autograph/converters/call_trees_test.py b/tensorflow/python/autograph/converters/call_trees_test.py index eb69d718a4c846..77c212188c9e30 100644 --- a/tensorflow/python/autograph/converters/call_trees_test.py +++ b/tensorflow/python/autograph/converters/call_trees_test.py @@ -14,6 +14,7 @@ # ============================================================================== """Tests for call_trees module.""" +import random import types from tensorflow.python.autograph.converters import call_trees @@ -257,6 +258,53 @@ def test_method(self, a): self.assertEqual(321, tr(tc, 1)) self.assertListEqual(mock.calls, [((1,), None)]) + def test_python_random_warning(self): + """Test that using Python random module triggers a warning.""" + # Reset the warning flag for testing + call_trees.python_random_warned = False + + def f(): + return random.randint(1, 10) + + # Transform should work and issue warning + tr, mock = self._transform_with_mock(f) + + # The function should still be callable + result = tr() + self.assertIsInstance(result, int) + self.assertGreaterEqual(result, 1) + self.assertLessEqual(result, 10) + + def test_python_random_randrange_warning(self): + """Test that using Python random.randrange triggers a warning.""" + # Reset the warning flag for testing + call_trees.python_random_warned = False + + def f(): + return random.randrange(0, 100) + + tr, mock = self._transform_with_mock(f) + + # The function should still be callable + result = tr() + self.assertIsInstance(result, int) + self.assertGreaterEqual(result, 0) + self.assertLess(result, 100) + + def test_python_random_choice_warning(self): + """Test that using Python random.choice triggers a warning.""" + # Reset the warning flag for testing + call_trees.python_random_warned = False + + def f(): + return random.choice([1, 2, 3]) + + tr, mock = self._transform_with_mock(f) + + # The function should still be callable + result = tr() + self.assertIn(result, [1, 2, 3]) + if __name__ == '__main__': test.main() From 14ac9d6ac831ab1c32304da774737085402f57f6 Mon Sep 17 00:00:00 2001 From: Saksham Singh Rathore Date: Thu, 30 Jul 2026 19:27:25 +0530 Subject: [PATCH 2/4] Fix Python random warning in AutoGraph --- .../python/autograph/converters/call_trees.py | 2 +- .../autograph/converters/call_trees_test.py | 22 +++++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tensorflow/python/autograph/converters/call_trees.py b/tensorflow/python/autograph/converters/call_trees.py index 7ebdc4b40b8ba9..6a1a584a6bcbd2 100644 --- a/tensorflow/python/autograph/converters/call_trees.py +++ b/tensorflow/python/autograph/converters/call_trees.py @@ -227,7 +227,7 @@ def visit_Call(self, node): 'unexpected behavior, especially with XLA compilation. ' 'Use `tf.random` functions instead for dynamic random values. ' 'For example, replace `random.randint(a, b)` with ' - '`tf.random.uniform([], a, b, dtype=tf.int32)`. ' + '`tf.random.uniform([], a, b + 1, dtype=tf.int32)`. ' 'See https://www.tensorflow.org/guide/function#executing_python_side_effects', full_name) python_random_warned = True diff --git a/tensorflow/python/autograph/converters/call_trees_test.py b/tensorflow/python/autograph/converters/call_trees_test.py index 77c212188c9e30..be1c0b7ca65bd2 100644 --- a/tensorflow/python/autograph/converters/call_trees_test.py +++ b/tensorflow/python/autograph/converters/call_trees_test.py @@ -266,8 +266,12 @@ def test_python_random_warning(self): def f(): return random.randint(1, 10) - # Transform should work and issue warning - tr, mock = self._transform_with_mock(f) + with self.assertLogs(level='WARNING') as logs: + tr, mock = self._transform_with_mock(f) + + self.assertLen(logs.output, 1) + self.assertIn('Detected use of Python\'s `random.randint()` inside a tf.function.', + logs.output[0]) # The function should still be callable result = tr() @@ -283,7 +287,12 @@ def test_python_random_randrange_warning(self): def f(): return random.randrange(0, 100) - tr, mock = self._transform_with_mock(f) + with self.assertLogs(level='WARNING') as logs: + tr, mock = self._transform_with_mock(f) + + self.assertLen(logs.output, 1) + self.assertIn('Detected use of Python\'s `random.randrange()` inside a tf.function.', + logs.output[0]) # The function should still be callable result = tr() @@ -299,7 +308,12 @@ def test_python_random_choice_warning(self): def f(): return random.choice([1, 2, 3]) - tr, mock = self._transform_with_mock(f) + with self.assertLogs(level='WARNING') as logs: + tr, mock = self._transform_with_mock(f) + + self.assertLen(logs.output, 1) + self.assertIn('Detected use of Python\'s `random.choice()` inside a tf.function.', + logs.output[0]) # The function should still be callable result = tr() From 69dde8efa13d48304b08d5c96c72bbc2a9ea7a98 Mon Sep 17 00:00:00 2001 From: Saksham Singh Rathore Date: Thu, 30 Jul 2026 19:35:33 +0530 Subject: [PATCH 3/4] Wrap AutoGraph random warning assertions --- .../autograph/converters/call_trees_test.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tensorflow/python/autograph/converters/call_trees_test.py b/tensorflow/python/autograph/converters/call_trees_test.py index be1c0b7ca65bd2..d8f90ff1c359ad 100644 --- a/tensorflow/python/autograph/converters/call_trees_test.py +++ b/tensorflow/python/autograph/converters/call_trees_test.py @@ -270,8 +270,9 @@ def f(): tr, mock = self._transform_with_mock(f) self.assertLen(logs.output, 1) - self.assertIn('Detected use of Python\'s `random.randint()` inside a tf.function.', - logs.output[0]) + self.assertIn( + 'Detected use of Python\'s `random.randint()` inside a tf.function.', + logs.output[0]) # The function should still be callable result = tr() @@ -291,8 +292,9 @@ def f(): tr, mock = self._transform_with_mock(f) self.assertLen(logs.output, 1) - self.assertIn('Detected use of Python\'s `random.randrange()` inside a tf.function.', - logs.output[0]) + self.assertIn( + 'Detected use of Python\'s `random.randrange()` inside a tf.function.', + logs.output[0]) # The function should still be callable result = tr() @@ -312,8 +314,9 @@ def f(): tr, mock = self._transform_with_mock(f) self.assertLen(logs.output, 1) - self.assertIn('Detected use of Python\'s `random.choice()` inside a tf.function.', - logs.output[0]) + self.assertIn( + 'Detected use of Python\'s `random.choice()` inside a tf.function.', + logs.output[0]) # The function should still be callable result = tr() From fb9f3d914d274e5b6daf1b31fa8c5bf0052092e1 Mon Sep 17 00:00:00 2001 From: jmestwa-coder Date: Sat, 29 Aug 2026 14:42:09 -0700 Subject: [PATCH 4/4] PR #118856: Validate metadata-driven chunk indices before vector access Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/118856 ## Summary Adds centralized validation for metadata-driven `chunk_index` values before vector access in proto splitter merge and fingerprint parsing paths. Malformed `.cpb` metadata could previously trigger unchecked indexing into `chunks` / `chunks_info` vectors during merge and fingerprint processing. This change introduces a shared `ValidateChunkIndex()` helper and applies it consistently across all affected metadata-driven lookup paths. ## Changes - Add shared `ValidateChunkIndex(uint64_t, size_t)` helper in `proto_splitter/cc/util.{h,cc}` - Validate root chunk indices in `Merger::Merge()` - Validate root chunk indices in `Merger::ReadFields()` - Validate nested chunk indices in `Merger::ProcessField()` (`READ` and `MERGE`) - Validate chunk indices in SavedModel fingerprint parsing paths - Add regression coverage for invalid root and nested chunk indices - Add regression coverage for invalid fingerprint metadata ## Behavior Before: - Malformed metadata could reach unchecked vector indexing paths After: - Invalid chunk indices fail deterministically with `absl::StatusCode::kFailedPrecondition` ## Testing Added regression tests for: - invalid root merge chunk index - invalid nested merge chunk index - invalid `ReadPartial()` metadata - invalid fingerprint parsing chunk index Copybara import of the project: -- 575231f30de391bfe42cd5bed25382fd5af066ae by jmestwa-coder : Validate metadata-driven chunk indices before vector access Merging this change closes #118856 COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/118856 from jmestwa-coder:validate-chunk-index-bounds 575231f30de391bfe42cd5bed25382fd5af066ae PiperOrigin-RevId: 973205398 --- .../cc/saved_model/fingerprinting_utils.cc | 4 ++ .../saved_model/fingerprinting_utils_test.cc | 35 ++++++++++ tensorflow/tools/proto_splitter/cc/util.cc | 7 ++ tensorflow/tools/proto_splitter/cc/util.h | 4 ++ tensorflow/tools/proto_splitter/merge.cc | 8 +++ tensorflow/tools/proto_splitter/merge_test.cc | 66 +++++++++++++++++++ 6 files changed, 124 insertions(+) diff --git a/tensorflow/cc/saved_model/fingerprinting_utils.cc b/tensorflow/cc/saved_model/fingerprinting_utils.cc index 218128bd60b4b6..92e61671c9132d 100644 --- a/tensorflow/cc/saved_model/fingerprinting_utils.cc +++ b/tensorflow/cc/saved_model/fingerprinting_utils.cc @@ -219,6 +219,8 @@ absl::StatusOr HashFields( if (chunked_message.has_chunk_index() && matches == field_tags.size()) { // chunked_field_tags are an exact match with field_tags. Hash referenced // chunk. + TF_RETURN_IF_ERROR(tools::proto_splitter::ValidateChunkIndex( + chunked_message.chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( std::string chunk, ReadChunk(reader, chunks_info[chunked_message.chunk_index()])); @@ -244,6 +246,8 @@ absl::StatusOr HashFields( merged_message = mfr.parent->GetReflection()->MutableMessage(mfr.parent, mfr.field); } + TF_RETURN_IF_ERROR(tools::proto_splitter::ValidateChunkIndex( + chunked_message.chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( std::string chunk, ReadChunk(reader, chunks_info[chunked_message.chunk_index()])); diff --git a/tensorflow/cc/saved_model/fingerprinting_utils_test.cc b/tensorflow/cc/saved_model/fingerprinting_utils_test.cc index e457ef128dcd3f..d51c01682c5fec 100644 --- a/tensorflow/cc/saved_model/fingerprinting_utils_test.cc +++ b/tensorflow/cc/saved_model/fingerprinting_utils_test.cc @@ -416,6 +416,41 @@ TEST(FingerprintingTest, CreateFingerprintDefNonExistentDirectoryReturnsError) { EXPECT_FALSE(result.ok()); } +TEST(FingerprintingTest, TestHashFieldsReturnsErrorOnInvalidChunkIndex) { + std::string cpb_file = io::JoinPath( + TensorFlowSrcRoot(), "tools/proto_splitter/testdata", "many-field.cpb"); + TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file)); + + auto read_metadata = GetChunkMetadata(reader); + if (!read_metadata.ok()) { + reader.Close(); + TF_ASSERT_OK(read_metadata.status()); + } + ChunkMetadata chunk_metadata = read_metadata.value(); + + ChunkedMessage invalid_chunked_message; + invalid_chunked_message.set_chunk_index(0); + auto* invalid_field = invalid_chunked_message.add_chunked_fields(); + invalid_field->add_field_tag()->set_field(1); + + std::vector chunks_info = std::vector( + chunk_metadata.chunks().begin(), chunk_metadata.chunks().end()); + invalid_field->mutable_message()->set_chunk_index(chunks_info.size()); + + RepeatedPtrField target_tags; + target_tags.Add()->set_field(1); + + ManyFields merged_message; + + auto statusor = HashFields(invalid_chunked_message, reader, chunks_info, + target_tags, &merged_message); + + EXPECT_FALSE(statusor.ok()); + EXPECT_EQ(statusor.status().code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(statusor.status().message()), + ::testing::HasSubstr("is out of range")); +} + } // namespace } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/tools/proto_splitter/cc/util.cc b/tensorflow/tools/proto_splitter/cc/util.cc index b38cdea137d444..095e6342d64879 100644 --- a/tensorflow/tools/proto_splitter/cc/util.cc +++ b/tensorflow/tools/proto_splitter/cc/util.cc @@ -801,5 +801,12 @@ absl::StatusOr OnlyContainsPb(absl::string_view prefix) { return false; } +absl::Status ValidateChunkIndex(uint64_t chunk_index, size_t chunks_size) { + if (chunk_index < chunks_size) return absl::OkStatus(); + return absl::FailedPreconditionError(absl::StrCat("Chunk index ", chunk_index, + " is out of range for ", + chunks_size, " chunks.")); +} + } // namespace tools::proto_splitter } // namespace tensorflow diff --git a/tensorflow/tools/proto_splitter/cc/util.h b/tensorflow/tools/proto_splitter/cc/util.h index ee525d752acf6c..54eff47aa047a4 100644 --- a/tensorflow/tools/proto_splitter/cc/util.h +++ b/tensorflow/tools/proto_splitter/cc/util.h @@ -15,6 +15,7 @@ limitations under the License. #ifndef TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_ #define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_ +#include #include #include #include @@ -164,6 +165,9 @@ absl::StatusOr ReadChunk( // file exists. Returns an error if neither .pb nor .cpb exist. absl::StatusOr OnlyContainsPb(absl::string_view prefix); +// Validates that a chunk index is within the bounds of a chunks vector. +absl::Status ValidateChunkIndex(uint64_t chunk_index, size_t chunks_size); + } // namespace tools::proto_splitter } // namespace tensorflow diff --git a/tensorflow/tools/proto_splitter/merge.cc b/tensorflow/tools/proto_splitter/merge.cc index d2243312769ce9..514e3f33caf810 100644 --- a/tensorflow/tools/proto_splitter/merge.cc +++ b/tensorflow/tools/proto_splitter/merge.cc @@ -59,6 +59,8 @@ absl::Status Merger::Merge(const std::vector>& chunks, if (chunked_message.has_chunk_index()) { // Chunks referenced by fields should be merged into the parent chunk. + TF_RETURN_IF_ERROR( + ValidateChunkIndex(chunked_message.chunk_index(), chunks.size())); merged_message->MergeFrom(*chunks[chunked_message.chunk_index()].get()); } @@ -237,6 +239,8 @@ absl::Status Merger::ReadFields(const ChunkedMessage& chunked_message, tsl::protobuf::Message* merged_message) { if (chunked_message.has_chunk_index()) { // Chunks referenced by fields should be merged into the parent chunk. + TF_RETURN_IF_ERROR( + ValidateChunkIndex(chunked_message.chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( std::string chunk, ReadChunk(reader, chunks_info[chunked_message.chunk_index()])); @@ -305,12 +309,16 @@ absl::Status Merger::ProcessField( std::string chunk; switch (op) { case MergerOp::READ: { + TF_RETURN_IF_ERROR(ValidateChunkIndex( + chunked_field.message().chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( chunk, ReadChunk(reader, chunks_info[chunked_field.message().chunk_index()])); break; } case MergerOp::MERGE: { + TF_RETURN_IF_ERROR(ValidateChunkIndex( + chunked_field.message().chunk_index(), chunks.size())); chunk = chunks[chunked_field.message().chunk_index()]->SerializeAsString(); break; diff --git a/tensorflow/tools/proto_splitter/merge_test.cc b/tensorflow/tools/proto_splitter/merge_test.cc index b49c901c7f9745..f0f2c0c73f7828 100644 --- a/tensorflow/tools/proto_splitter/merge_test.cc +++ b/tensorflow/tools/proto_splitter/merge_test.cc @@ -312,6 +312,72 @@ TEST(MergeTest, TestProcessFieldReturnsErrorOnInvalidFieldNumber) { ::testing::HasSubstr("not found in message descriptor")); } +TEST(MergeTest, TestMergeReturnsErrorOnInvalidRootChunkIndex) { + ::tensorflow::proto_splitter::ChunkedMessage chunked_message; + chunked_message.set_chunk_index(1); + + std::vector> chunks; + chunks.push_back( + std::make_unique<::tensorflow::proto_splitter_testdata::ManyFields>()); + + ::tensorflow::proto_splitter_testdata::ManyFields merged; + absl::Status status = Merger::Merge(chunks, chunked_message, &merged); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(status.message()), + ::testing::HasSubstr("Chunk index 1 is out of range")); +} + +TEST(MergeTest, TestMergeReturnsErrorOnInvalidNestedChunkIndex) { + ::tensorflow::proto_splitter::ChunkedMessage chunked_message; + auto* chunk_field = chunked_message.add_chunked_fields(); + auto* tag = chunk_field->add_field_tag(); + tag->set_field(1); + chunk_field->mutable_message()->set_chunk_index(1); + + std::vector> chunks; + chunks.push_back( + std::make_unique<::tensorflow::proto_splitter_testdata::ManyFields>()); + + ::tensorflow::proto_splitter_testdata::ManyFields merged; + absl::Status status = Merger::Merge(chunks, chunked_message, &merged); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(status.message()), + ::testing::HasSubstr("Chunk index 1 is out of range")); +} + +TEST(MergeTest, TestReadPartialReturnsErrorOnInvalidRootChunkIndex) { + const std::string path = + io::JoinPath(testing::TensorFlowSrcRoot(), + "tools/proto_splitter/testdata", "many-field"); + TF_ASSERT_OK_AND_ASSIGN(auto reader, tools::proto_splitter::GetRiegeliReader( + absl::StrCat(path, ".cpb"))); + + auto read_metadata = GetChunkMetadata(reader); + if (!read_metadata.ok()) { + reader.Close(); + TF_ASSERT_OK(read_metadata.status()); + } + reader.Close(); + + ::tensorflow::proto_splitter::ChunkMetadata chunk_metadata = + read_metadata.value(); + chunk_metadata.mutable_message()->set_chunk_index( + chunk_metadata.chunks_size()); + + proto_splitter_testdata::ManyFields merged_many_fields; + absl::Status status = + Merger::ReadPartial(path, chunk_metadata, &merged_many_fields); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(status.message()), + ::testing::HasSubstr("is out of range")); +} + } // namespace } // namespace tensorflow::tools::proto_splitter