code: def testConv2DBatching(self, layer_batch, input_batch):
      dist = jdlayers.Conv2D(4, (3, 3), 3, dtype=self.dtype)
      layer = dist.sample(
          layer_batch,
          seed=test_util.test_seed(sampler_type="stateless"),
      )
      # Validate that we can map the layer.
      layer = tf.nest.map_structure(lambda x: x + 0.0, layer)
      x = tf.ones(input_batch + [9, 5, 5, 3], dtype=self.dtype)
      y = layer(x)
      self.assertAllEqual(
          list(ps.broadcast_shape(layer_batch, input_batch))
          + [9, 5, 5, 4],
          y.shape,
      )
      self.assertEqual(self.dtype, y.dtype)
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import real_nvp
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.experimental.joint_distribution_layers import layers
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 5 for calling the _JDLayersTestBase obj's underlying member method code with example usage: obj.testConv2DBatching(layer_batch, input_batch) without return types