code: def make_convolution_fn(
      filter_shape,
      rank,
      strides,
      padding,
      dilations=None,
      dtype=tf.int32,
      validate_args=False,
      name=None,
  ):
      """Like `tf.nn.conv2d` except applies batch of kernels to batch of `x`."""
      with tf.name_scope(name or "conv2d"):
          if tf.get_static_value(rank) != 2:
              raise NotImplementedError(
                  "Argument `rank` currently only supports `2`; "
                  'saw "{}".'.format(rank)
              )
          [
              filter_shape,
              rank,
              strides,
              padding,
              dilations,
          ] = prepare_conv_args(
              filter_shape,
              rank=rank,
              strides=strides,
              padding=padding,
              dilations=dilations,
              validate_args=validate_args,
          )
  
      def op(x, kernel):
          input_dtype = dtype_util.common_dtype(
              [x, kernel], dtype_hint=tf.float32
          )
          x = tf.convert_to_tensor(x, dtype=input_dtype, name="x")
          kernel = tf.convert_to_tensor(
              kernel, dtype=input_dtype, name="kernel"
          )
  
          batch_shape, event_shape = ps.split(
              ps.shape(x), num_or_size_splits=[-1, 3]
          )
          xh, xw, c_in = ps.unstack(event_shape, num=3)
          fh, fw = filter_shape
  
          assertions = _maybe_validate_input_shapes(
              ps.shape(kernel),
              channels_in=c_in,
              filter_height=fh,
              filter_width=fw,
              validate_args=validate_args,
          )
  
          with tf.control_dependencies(assertions):
              if tf.get_static_value(ps.rank(kernel)) == 2:
                  flat_x = tf.reshape(
                      x, shape=ps.concat([[-1], event_shape], axis=0)
                  )
                  flat_y = tf.nn.conv2d(
                      x,
                      filters=tf.reshape(
                          kernel, shape=[fh, fw, c_in, -1]
                      ),
                      strides=strides,
                      padding=padding,
                      data_format="NHWC",
                      dilations=dilations,
                  )
                  output_shape = ps.shape(flat_y)[-3:]
                  return tf.reshape(
                      flat_y,
                      shape=ps.concat(
                          [batch_shape, output_shape], axis=0
                      ),
                  )
  
              pad_values = [
                  _get_conv_padding(
                      xdim,
                      filter_dim=k,
                      stride=s,
                      dilation=d,
                      padding=padding,
                  )
                  for (xdim, k, s, d) in zip(
                      (xh, xw), filter_shape, strides, dilations
                  )
              ]
  
              idx, shape = im2row_index(
                  (
                      xh + sum(pad_values[0]),
                      xw + sum(pad_values[1]),
                      c_in,
                  ),
                  block_shape=filter_shape,
                  slice_step=strides,
                  dilations=dilations,
                  dtype=dtype,
              )
  
              if padding == "SAME":
                  n = ps.maximum(0, ps.rank(x) - 3)
                  paddings = ps.pad(
                      pad_values,
                      paddings=[[n, 1], [0, 0]],
                      constant_values=0,
                  )
                  x = tf.pad(x, paddings=paddings, constant_values=0)
  
              flat_shape = ps.pad(
                  batch_shape, paddings=[[0, 1]], constant_values=-1
              )
              flat_x = tf.gather(
                  tf.reshape(x, shape=flat_shape), indices=idx, axis=-1
              )
              im_x = tf.reshape(
                  flat_x, shape=ps.concat([batch_shape, shape], axis=0)
              )
              return tf.matmul(im_x, kernel[..., tf.newaxis, :, :])
  
      return op
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.nn.util import utils
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 5 for calling a global function with example usage: make_convolution_fn(filter_shape, rank, strides, padding, dilations, dtype, validate_args, name) and returns: op, tf, tf