code: def testNumParallelBatches(self, verify_fn, drop_remainder):
      range_size = 11
      num_shards = 3
      num_repeats = 2
      batch_size = 5
      num_parallel_batches = 2
      total_outputs = (range_size // num_shards) * num_repeats
      if drop_remainder:
          num_outputs = total_outputs // batch_size
      else:
          num_outputs = int(math.ceil(total_outputs / batch_size))
  
      def build_ds(range_start, drop_remainder):
          def _map_fn(x):
              return math_ops.square(x)
  
          return (
              dataset_ops.Dataset.range(
                  range_start, range_start + range_size
              )
              .shard(num_shards=num_shards, index=0)
              .repeat(num_repeats)
              .apply(
                  batching.map_and_batch(
                      map_func=_map_fn,
                      batch_size=batch_size,
                      num_parallel_batches=num_parallel_batches,
                      drop_remainder=drop_remainder,
                  )
              )
          )
  
      verify_fn(
          self,
          lambda: build_ds(10, drop_remainder=drop_remainder),
          num_outputs,
      )
  
imports: import math
  from absl.testing import parameterized
  import numpy
  from tensorflow.python import pywrap_sanitizers
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.data.experimental.ops import batching
  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.eager import context
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import constant_op
  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 cond
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import script_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling the MapAndBatchCheckpointTest obj's underlying member method code with example usage: obj.testNumParallelBatches(verify_fn, drop_remainder) and returns: dataset_ops, math_ops