tf-models-nightly 2.17.0.dev20240617__py2.py3-none-any.whl → 2.20.0.dev20251220__py2.py3-none-any.whl

official/vision/serving/export_tflite.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_tflite_lib.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/export_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/image_classification.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/image_classification_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/semantic_segmentation.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -26,7 +26,8 @@ class SegmentationModule(export_base.ExportModule):
 
   def _build_model(self):
     input_specs = tf_keras.layers.InputSpec(
-        shape=[self._batch_size] + self._input_image_size + [3])
+        shape=[self._batch_size] + self._input_image_size + [self._num_channels]
+    )
 
     return factory.build_segmentation_model(
         input_specs=input_specs,
@@ -72,7 +73,9 @@ class SegmentationModule(export_base.ExportModule):
     if self._input_type != 'tflite':
       with tf.device('cpu:0'):
         images_spec = tf.TensorSpec(
-            shape=self._input_image_size + [3], dtype=tf.float32)
+            shape=self._input_image_size + [self._num_channels],
+            dtype=tf.float32,
+        )
         image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
 
         images, image_info = tf.nest.map_structure(

official/vision/serving/semantic_segmentation_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -49,20 +49,22 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
         {input_type: 'serving_default'})
     tf.saved_model.save(module, save_directory, signatures=signatures)
 
-  def _get_dummy_input(self, input_type, input_image_size):
+  def _get_dummy_input(self, input_type, input_image_size, num_channels):
     """Get dummy input for the given input type."""
 
     height = input_image_size[0]
     width = input_image_size[1]
     if input_type == 'image_tensor':
-      return tf.zeros((1, height, width, 3), dtype=np.uint8)
+      return tf.zeros((1, height, width, num_channels), dtype=np.uint8)
     elif input_type == 'image_bytes':
-      image = Image.fromarray(np.zeros((height, width, 3), dtype=np.uint8))
+      image = Image.fromarray(
+          np.zeros((height, width, num_channels), dtype=np.uint8)
+      )
       byte_io = io.BytesIO()
       image.save(byte_io, 'PNG')
       return [byte_io.getvalue()]
     elif input_type == 'tf_example':
-      image_tensor = tf.zeros((height, width, 3), dtype=tf.uint8)
+      image_tensor = tf.zeros((height, width, num_channels), dtype=tf.uint8)
       encoded_jpeg = tf.image.encode_jpeg(tf.constant(image_tensor)).numpy()
       example = tf.train.Example(
           features=tf.train.Features(
@@ -73,7 +75,7 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
               })).SerializeToString()
       return [example]
     elif input_type == 'tflite':
-      return tf.zeros((1, height, width, 3), dtype=np.float32)
+      return tf.zeros((1, height, width, num_channels), dtype=np.float32)
 
   @parameterized.parameters(
       ('image_tensor', False, [112, 112], False),
@@ -105,7 +107,7 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
     imported = tf.saved_model.load(tmp_dir)
     segmentation_fn = imported.signatures['serving_default']
 
-    images = self._get_dummy_input(input_type, input_image_size)
+    images = self._get_dummy_input(input_type, input_image_size, num_channels=3)
     if input_type != 'tflite':
       processed_images, _ = tf.nest.map_structure(
           tf.stop_gradient,
@@ -128,6 +130,68 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
     out = segmentation_fn(tf.constant(images))
     self.assertAllClose(out['logits'].numpy(), expected_output.numpy())
 
+  @parameterized.parameters(
+      ('image_tensor',),
+      ('tflite',),
+  )
+  def test_export_with_extra_input_channels(self, input_type):
+    tmp_dir = self.get_temp_dir()
+    num_channels = 6
+    params = exp_factory.get_exp_config('mnv2_deeplabv3_pascal')
+    params.task.init_checkpoint = None
+    params.task.model.input_size = [112, 112, num_channels]
+    params.task.export_config.rescale_output = False
+    params.task.train_data.preserve_aspect_ratio = False
+    params.task.train_data.image_feature.mean = [0.5] * num_channels
+    params.task.train_data.image_feature.stddev = [0.5] * num_channels
+    params.task.train_data.image_feature.num_channels = num_channels
+    module = semantic_segmentation.SegmentationModule(
+        params,
+        batch_size=1,
+        input_image_size=[112, 112],
+        input_type=input_type,
+        num_channels=num_channels,
+    )
+
+    self._export_from_module(module, input_type, tmp_dir)
+
+    self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
+    self.assertTrue(
+        os.path.exists(os.path.join(tmp_dir, 'variables', 'variables.index'))
+    )
+    self.assertTrue(
+        os.path.exists(
+            os.path.join(tmp_dir, 'variables', 'variables.data-00000-of-00001')
+        )
+    )
+
+    imported = tf.saved_model.load(tmp_dir)
+    segmentation_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(input_type, [112, 112], num_channels)
+
+    if input_type != 'tflite':
+      processed_images, _ = tf.nest.map_structure(
+          tf.stop_gradient,
+          tf.map_fn(
+              module._build_inputs,
+              elems=tf.zeros((1, 112, 112, num_channels), dtype=tf.uint8),
+              fn_output_signature=(
+                  tf.TensorSpec(
+                      shape=[112, 112, num_channels], dtype=tf.float32
+                  ),
+                  tf.TensorSpec(shape=[4, 2], dtype=tf.float32),
+              ),
+          ),
+      )
+    else:
+      processed_images = images
+
+    logits = module.model(processed_images, training=False)['logits']
+    expected_output = tf.image.resize(logits, [112, 112], method='bilinear')
+    out = segmentation_fn(tf.constant(images))
+    self.assertAllClose(out['logits'].numpy(), expected_output.numpy())
+
   def test_export_invalid_batch_size(self):
     batch_size = 3
     tmp_dir = self.get_temp_dir()

official/vision/serving/video_classification.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/serving/video_classification_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/tasks/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/tasks/image_classification.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/tasks/maskrcnn.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/tasks/retinanet.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/tasks/semantic_segmentation.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/tasks/video_classification.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/train.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/train_spatial_partitioning.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/argmax_matcher.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/balanced_positive_negative_sampler.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/box_coder.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/box_list.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/box_list_ops.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/faster_rcnn_box_coder.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/matcher.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/minibatch_sampler.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/ops.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/preprocessor.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/region_similarity_calculator.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/shape_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/target_assigner.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/object_detection/visualization_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -438,6 +438,11 @@ def visualize_outputs(
       )
     return tf.cast(images, dtype=tf.uint8)
 
+  if images.shape[3] > 3:
+    images = images[:, :, :, 0:3]
+  elif images.shape[3] == 1:
+    images = tf.image.grayscale_to_rgb(images)
+
   images = tf.nest.map_structure(
       tf.identity,
       tf.map_fn(

official/vision/utils/ops_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/utils/summary_manager.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/actions/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/actions/conditional_action.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -47,7 +47,8 @@ class ConditionalAction:
     """Initializes the instance.
 
     Args:
-      condition: A callable accepting train or eval outputs and returing a bool.
+      condition: A callable accepting train or eval outputs and returning
+       a bool.
       action: The action (or optionally sequence of actions) to perform when
         `condition` is met.
     """

orbit/actions/conditional_action_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/actions/export_saved_model.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/actions/export_saved_model_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/actions/new_best_metric.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -40,7 +40,7 @@ class NewBestMetric:
   if it is achieved. These separate methods enable the same `NewBestMetric`
   instance to be reused as a condition multiple times, and can also provide
   additional preemption/failure safety. For example, to avoid updating the best
-  metric if a model export fails or is pre-empted:
+  metric if a model export fails or is pre-emptied:
 
       new_best_metric = orbit.actions.NewBestMetric(
         'accuracy', filename='/model/dir/best_metric')

orbit/actions/new_best_metric_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -61,7 +61,7 @@ class NewBestMetricTest(tf.test.TestCase):
     tempfile = self.create_tempfile().full_path
     value = {'a': 1, 'b': 2}
     persisted_value = actions.JSONPersistedValue(value, tempfile)
-    # The inital value is used since tempfile is empty.
+    # The initial value is used since tempfile is empty.
     self.assertEqual(persisted_value.read(), value)
     persisted_value = actions.JSONPersistedValue('ignored', tempfile)
     # Initial value of 'ignored' is ignored, since there's a value in tempfile.

orbit/actions/save_checkpoint_if_preempted.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/controller.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/controller_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/examples/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/examples/single_task/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/examples/single_task/single_task_evaluator.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/examples/single_task/single_task_evaluator_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/examples/single_task/single_task_trainer.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/examples/single_task/single_task_trainer_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/runner.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/standard_runner.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/standard_runner_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/utils/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/utils/common.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/utils/common_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/utils/epoch_helper.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

orbit/utils/loop_fns.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Orbit Authors. All Rights Reserved.
+# Copyright 2025 The Orbit Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,6 +19,10 @@ from orbit.utils import tpu_summaries
 
 import tensorflow as tf, tf_keras
 
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.tpu import embedding_context_utils as ecu
+# pylint: enable=g-direct-tensorflow-import
+
 
 def create_loop_fn(step_fn):
   """Creates a loop function driven by a Python `while` loop.
@@ -200,7 +204,8 @@ class LoopFnWithSummaries(tpu_summaries.OptionalSummariesFunction):
 
   def __call__(self, iterator, num_steps):
     if tf.summary.should_record_summaries():
-      output = self.with_summaries(iterator, tf.constant(1))
+      with ecu.SequentialEmbeddingContext():
+        output = self.with_summaries(iterator, tf.constant(1))
       num_steps -= 1
     if num_steps >= 1:
       output = self.without_summaries(iterator, num_steps)