class_docstr: Test cases for TfRecordRepresentativeDatasetLoader.
class_docstr_tok: ["Test","cases","for","TfRecordRepresentativeDatasetLoader","."]
code: def test_tf_record_saver_with_generator_dataset(self):
      tf_record_path = self.create_tempfile().full_path
      path_map = {"serving_default": tf_record_path}
      num_samples = 2
  
      def data_gen():
          for _ in range(num_samples):
              yield {"x": [1, 2]}
  
      repr_ds_map = {"serving_default": data_gen()}
      saver = repr_dataset.TfRecordRepresentativeDatasetSaver(path_map)
      dataset_file_map = saver.save(repr_ds_map)
      self.assertCountEqual(
          dataset_file_map.keys(), ["serving_default"]
      )
  
      dataset_map = repr_dataset.TfRecordRepresentativeDatasetLoader(
          dataset_file_map
      ).load()
      self.assertCountEqual(dataset_map.keys(), ["serving_default"])
      samples = dataset_map["serving_default"]
      for sample in samples:
          self.assertCountEqual(sample.keys(), {"x"})
          self.assertAllEqual(sample["x"], np.array([1, 2]))
  
      self.assertLen(samples, num_samples)
  
imports: import random
  import numpy
  from tensorflow.compiler.mlir.quantization.tensorflow.python import representative_dataset
  from tensorflow.core.protobuf import meta_graph_pb2
  from tensorflow.python.client import session
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.platform import test
  from tensorflow.python.types import core
label_desc: Use image node_id 1 for calling the TfRecordRepresentativeDatasetTest obj's underlying member method code with example usage: obj.test_tf_record_saver_with_generator_dataset() without return types