code: def window_partition(x, window_size: int):
      """
      将feature map按照window_size划分成一个个没有重叠的window
      Args:
          x: (B, H, W, C)
          window_size (int): window size(M)
  
      Returns:
          windows: (num_windows*B, window_size, window_size, C)
      """
      B, H, W, C = x.shape
      x = tf.reshape(
          x,
          [
              B,
              H // window_size,
              window_size,
              W // window_size,
              window_size,
              C,
          ],
      )
      # transpose: [B, H//Mh, Mh, W//Mw, Mw, C] -> [B, H//Mh, W//Mh, Mw, Mw, C]
      # reshape: [B, H//Mh, W//Mw, Mh, Mw, C] -> [B*num_windows, Mh, Mw, C]
      x = tf.transpose(x, [0, 1, 3, 2, 4, 5])
      windows = tf.reshape(x, [-1, window_size, window_size, C])
      return windows
  
imports: import tensorflow
  from tensorflow.keras import Model, layers, initializers
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: window_partition(x, window_size) and returns: windows