code: def test_batching(self, input_batch_shape, kernel_batch_shape):
      input_shape = (12, 12, 2)
      filter_shape = (2, 2)
      channels_out = 3
      strides = (1, 1)
      dilations = (1, 1)
      padding = "SAME"
  
      x, k = _make_input_and_kernel(
          self.make_input,
          input_batch_shape=input_batch_shape,
          input_shape=input_shape,
          kernel_batch_shape=kernel_batch_shape,
          filter_shape=filter_shape,
          channels_out=channels_out,
          dtype=self.dtype,
      )
  
      conv_fn = convolution_util.make_convolution_fn(
          filter_shape,
          rank=2,
          strides=strides,
          padding=padding,
          dilations=dilations,
          validate_args=True,
      )
      y_batched = conv_fn(x, k)
  
      broadcast_batch_shape = ps.broadcast_shape(
          input_batch_shape, kernel_batch_shape
      )
      broadcasted_input = tf.broadcast_to(
          x,
          shape=ps.concat([broadcast_batch_shape, input_shape], axis=0),
      )
      broadcasted_kernel = tf.broadcast_to(
          k,
          shape=ps.concat(
              [broadcast_batch_shape, ps.shape(k)[-2:]], axis=0
          ),
      )
  
      flat_y = tf.reshape(
          y_batched,
          shape=ps.pad(
              ps.shape(y_batched)[-3:],
              paddings=[[1, 0]],
              constant_values=-1,
          ),
      )
      flat_x = tf.reshape(
          broadcasted_input,
          shape=ps.pad(
              input_shape, paddings=[[1, 0]], constant_values=-1
          ),
      )
      flat_tf_kernel = tf.reshape(
          broadcasted_kernel,
          shape=ps.concat(
              [(-1,), filter_shape, (input_shape[-1], channels_out)],
              axis=0,
          ),
      )
  
      y_expected = tf.vectorized_map(
          lambda args: tf.nn.conv2d(  # pylint: disable=g-long-lambda
              args[0][tf.newaxis],
              args[1],
              strides=strides,
              padding=padding,
          ),
          elems=(flat_x, flat_tf_kernel),
      )
  
      [y_actual_, y_expected_] = self.evaluate(
          [flat_y, tf.squeeze(y_expected, axis=1)]
      )
      self.assertAllClose(y_expected_, y_actual_, rtol=1e-5, atol=0)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.nn.util import convolution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 2 for calling the _BatchedConvTest obj's underlying member method code with example usage: obj.test_batching(input_batch_shape, kernel_batch_shape) without return types