code: def im2row(
      x, block_shape, slice_step=(1, 1), padding="VALID", name=None
  ):
      """Rearrange image blocks into rows.
  
      This function can be used to implement 2D convolution as a `matmul`, e.g.,
  
      `tf.nn.conv2d(x, k) = tf.matmul(
          tf.experimental.nn.util.im2row(x), tf.reshape(k, shape=[-1, out_size]))`.
  
      Args:
        x: Rank 3 (or more) Tensor representing 2D images.
        block_shape: Length-2 vector representing the block or "filter" shape.
        slice_step: Length-2 vector specifying the convolution stride length.
          Default value: `(1, 1)`.
        padding: One of `'VALID'` or `'SAME'` (case insensitive).
          Default value: `'VALID'`.
        name: Python `str` used to describe ops created by this function.
          Default value: `None` (i.e., `'im2col'`).
  
      Returns:
        im2row_x: batch of matrices representing subblock copies of `x`.
          Same batch shape as `x` but with rightmost shape:
          `batch_shape + [oh * ow, block_shape[0] * block_shape[1] * channels]`,
          where `oh = (h - block_shape[0] + 1) // slice_step[0]` and
          `ow = (w - block_shape[1] + 1) // slice_step[1]` when `padding = 'VALID'`
          and `oh = h` and `ow = w` when `padding = 'SAME'`.
        shape: shape `Tensor` equivalent to:
          `batch_shape + [oh, ow, block_shape[0] * block_shape[1] * channels]` where
          `oh, ow` are defined as above.
      """
      with tf.name_scope(name or "im2row"):
          padding = _validate_padding(padding)
          if padding == "VALID":
              pass  # Do nothing.
          elif padding == "SAME":
              raise NotImplementedError(
                  'Argument padding="SAME" not implemented.'
              )
              # TODO(jvdillon): See if the following works:
              # fh, fw = block_shape
              # o = 1 if data_format == 'NHWC' else 0
              # n = ps.maximum(0, ps.rank(x) - 3)
              # paddings = ps.pad(
              #     [[0, fh - 1], [0, fw - 1]],
              #     paddings=[[n + 1 - o, o], [0, 0]],
              #     constant_values=0)
              # x = tf.pad(x, paddings=paddings, constant_values=0)
              # padding = 'VALID'
          else:
              assert False  # Can't be here.
          x_shape = ps.shape(x)
          idx, s = im2row_index(
              x_shape, block_shape=block_shape, slice_step=slice_step
          )
          flat_shape = ps.pad(
              x_shape[:-3], paddings=[[0, 1]], constant_values=-1
          )
          x = tf.gather(
              tf.reshape(x, flat_shape), idx, axis=-1
          )  # == np.take
          return tf.reshape(x, s)
  
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 1 for calling a global function with example usage: im2row(x, block_shape, slice_step, padding, name) and returns: tf