[pull] master from tensorflow:master

tensorflow/cc/saved_model/fingerprinting_utils.cc

@@ -219,6 +219,8 @@ absl::StatusOr<uint64_t> 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<uint64_t> 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()]));

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<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
+      chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
+  invalid_field->mutable_message()->set_chunk_index(chunks_info.size());
+
+  RepeatedPtrField<FieldIndex> 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

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,26 @@ 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 + 1, 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

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,73 @@ 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)
+
+    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()
+    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)
+
+    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()
+    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])
+
+    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()
+    self.assertIn(result, [1, 2, 3])
+
 
 if __name__ == '__main__':
   test.main()

tensorflow/tools/proto_splitter/cc/util.cc

@@ -801,5 +801,12 @@ absl::StatusOr<bool> 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

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 <cstddef>
 #include <cstdint>
 #include <functional>
 #include <memory>
@@ -164,6 +165,9 @@ absl::StatusOr<std::string> ReadChunk(
 // file exists. Returns an error if neither .pb nor .cpb exist.
 absl::StatusOr<bool> 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
 

tensorflow/tools/proto_splitter/merge.cc

@@ -59,6 +59,8 @@ absl::Status Merger::Merge(const std::vector<std::unique_ptr<Message>>& 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;

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<std::unique_ptr<tsl::protobuf::Message>> 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<std::unique_ptr<tsl::protobuf::Message>> 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