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
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"]
code: def testBasic(self, count, batch_size, drop_remainder):
      """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
      """
      dataset = dataset_ops.Dataset.from_tensor_slices(
          list(range(count))
      ).batch(batch_size=batch_size, drop_remainder=drop_remainder)
      num_full_batches = count // batch_size
      for i in range(num_full_batches):
          expected_batch = np.arange(
              i * batch_size,
              (i * batch_size + batch_size),
              1,
              dtype=np.int32,
          )
          self.assertAllEqual(
              expected_batch,
              self.evaluate(random_access.at(dataset, i)),
          )
      has_remainder = (not drop_remainder) and (count % batch_size != 0)
      if has_remainder:
          expected_batch = np.arange(
              batch_size * num_full_batches, count, 1
          )
          self.assertAllEqual(
              expected_batch,
              self.evaluate(
                  random_access.at(dataset, num_full_batches)
              ),
          )
      with self.assertRaises(errors.OutOfRangeError):
          self.evaluate(
              random_access.at(
                  dataset,
                  index=num_full_batches + (1 if has_remainder else 0),
              )
          )
  
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 3 for calling the BatchRandomAccessTest obj's underlying member method code with example usage: obj.testBasic(count, batch_size, drop_remainder) without return types