# coding=utf-8
# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for input_generator."""

from absl.testing import absltest
from lingvo.core import base_input_generator
from lingvo.core import test_utils
from paxml.tasks.vision import input_generator


class InputGeneratorTest(test_utils.TestCase):

  def _check_data_shape(self,
                        p,
                        is_multlabels=False,
                        batch_size=8,
                        image_size=33):
    p.data_shape = (image_size, image_size, 3)
    p.batch_size = batch_size
    if isinstance(p.cls, base_input_generator.BaseInputGeneratorFromFiles):
      p.num_batcher_threads = 1
      p.file_parallelism = 1
      p.file_buffer_size = 1
    with self.session() as sess:
      inp = p.Instantiate()
      batch = sess.run(inp.GetPreprocessedInputBatch())
      b, (h, w, d) = p.batch_size, p.data_shape
      self.assertEqual(batch.image.shape, (b, h, w, d))
      self.assertEqual(batch.label_probs.shape, (b, p.num_classes))

  def testImageNetValidation(self):
    p = input_generator.ImageNetValidation.Params()
    self._check_data_shape(p)

  def testImageNetTrain(self):
    p = input_generator.ImageNetTrain.Params()
    self._check_data_shape(p)


if __name__ == '__main__':
  absltest.main()

@@ -656,88 +656,8 @@
 est%0A
-from lingvo.core import base_input_generator%0Afrom lingvo.core import test_utils%0A
 from
@@ -730,18 +730,16 @@
 est(
+absl
 test
-_utils
 .Tes
@@ -1008,86 +1008,8 @@
 ize%0A
-    if isinstance(p.cls, base_input_generator.BaseInputGeneratorFromFiles):%0A  
     
@@ -1030,26 +1030,24 @@
 threads = 1%0A
-  
     p.file_p
@@ -1065,18 +1065,16 @@
 = 1%0A    
- 
- 
 p.file_b
@@ -1092,43 +1092,8 @@
 = 1%0A
-    with self.session() as sess:%0A  
     
@@ -1114,18 +1114,16 @@
 tiate()%0A
-  
     batc
@@ -1130,17 +1130,8 @@
 h = 
-sess.run(
 inp.
@@ -1157,20 +1157,17 @@
 tBatch()
-)%0A  
+%0A
     b, (
@@ -1204,26 +1204,24 @@
 a_shape%0A    
-  
 self.assertE
@@ -1258,18 +1258,16 @@
  w, d))%0A
-  
     self

Make vision input generator test case run in TF eager mode.