docstr: Tests the batch dataset logic for various input configurations.
  
  Args:
    count: the number of input elements
    batch_size: the batch size
    drop_remainder: whether a smaller batch size should be produced if batch
      size does not divide number of inputs evenly
    num_parallel_calls: the number batches to process asynchronously in
      parallel
docstr_tok: ["Tests","the","batch","dataset","logic","for","various","input","configurations",".","Args",":","count",":","the","number","of","input","elements","batch_size",":","the","batch","size","drop_remainder",":","whether","a","smaller","batch","size","should","be","produced","if","batch","size","does","not","divide","number","of","inputs","evenly","num_parallel_calls",":","the","number","batches","to","process","asynchronously","in","parallel"]
code: def testBasic(
      self, count, batch_size, drop_remainder, num_parallel_calls
  ):
      """Tests the batch dataset logic for various input configurations.
  
      Args:
        count: the number of input elements
        batch_size: the batch size
        drop_remainder: whether a smaller batch size should be produced if batch
          size does not divide number of inputs evenly
        num_parallel_calls: the number batches to process asynchronously in
          parallel
      """
  
      # The pipeline is TensorSliceDataset -> MapDataset(square_3) ->
      # RepeatDataset(count) -> BatchDataset(batch_size).
      components = (
          np.arange(7),
          np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
          np.array(37.0) * np.arange(7),
      )
  
      def _map_fn(x, y, z):
          return (
              math_ops.square(x),
              math_ops.square(y),
              math_ops.square(z),
          )
  
      dataset = (
          dataset_ops.Dataset.from_tensor_slices(components)
          .map(_map_fn)
          .repeat(count)
          .batch(batch_size, drop_remainder, num_parallel_calls)
      )
      get_next = self.getNext(dataset)
  
      if drop_remainder:
          dim0 = batch_size
      else:
          dim0 = None
      self.assertEqual(
          [
              ts.as_list()
              for ts in nest.flatten(
                  dataset_ops.get_legacy_output_shapes(dataset)
              )
          ],
          [[dim0] + list(c.shape[1:]) for c in components],
      )
  
      num_full_batches = (count * 7) // batch_size
      for i in range(num_full_batches):
          result = self.evaluate(get_next())
          for component, result_component in zip(components, result):
              for j in range(batch_size):
                  self.assertAllEqual(
                      component[(i * batch_size + j) % 7] ** 2,
                      result_component[j],
                  )
      if not drop_remainder and (count * 7) % batch_size > 0:
          result = self.evaluate(get_next())
          for component, result_component in zip(components, result):
              for j in range((count * 7) % batch_size):
                  self.assertAllEqual(
                      component[(num_full_batches * batch_size + j) % 7]
                      ** 2,
                      result_component[j],
                  )
      with self.assertRaises(errors.OutOfRangeError):
          result = self.evaluate(get_next())
  
imports: import time
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.data.experimental.ops import random_access
  from tensorflow.python.data.kernel_tests import checkpoint_test_base
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import options
  from tensorflow.python.data.util import nest
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import script_ops
  from tensorflow.python.ops.ragged import ragged_concat_ops
  from tensorflow.python.ops.ragged import ragged_factory_ops
  from tensorflow.python.ops.ragged import ragged_math_ops
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.platform import test
label_desc: Use image node_id 1 for calling the BatchTest obj's underlying member method code with example usage: obj.testBasic(count, batch_size, drop_remainder, num_parallel_calls) and returns: math_ops,  math_ops,  math_ops