code: def testShardCounts(self, mesh_dims, layout_specs, batch_dim, counts):
      num_devices = np.prod([size for _, size in mesh_dims])
      mesh = mesh_util.create_mesh(
          mesh_dims=mesh_dims,
          devices=["CPU:%d" % i for i in range(num_devices)],
      )
      layout = Layout(layout_specs, mesh)
  
      self.assertEqual(
          input_util._shard_counts(layout, batch_dim), counts
      )
  
imports: import contextlib
  import threading
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import mesh_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 1 for calling the InputUtilHelpersTest obj's underlying member method code with example usage: obj.testShardCounts(mesh_dims, layout_specs, batch_dim, counts) without return types