code: def blockify(x, block_size: int):
      """image to blocks
      Args:
          x (Tensor): with shape (B, H, W, C)
          block_size (int): edge length of a single square block in units of H, W
      """
      B, H, W, C = x.shape
      _assert(
          H % block_size == 0,
          "`block_size` must divide input height evenly",
      )
      _assert(
          W % block_size == 0,
          "`block_size` must divide input width evenly",
      )
      grid_height = H // block_size
      grid_width = W // block_size
      x = x.reshape(
          B, grid_height, block_size, grid_width, block_size, C
      )
      x = x.transpose(2, 3).reshape(B, grid_height * grid_width, -1, C)
      return x
  
imports: import collections.abc
  import logging
  import math
  from functools import partial
  import torch
  import torch.nn.functional
  from torch import nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import PatchEmbed, Mlp, DropPath, create_classifier, trunc_normal_, _assert
  from timm.layers import create_conv2d, create_pool2d, to_ntuple, use_fused_attn, LayerNorm
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_function
  from ._manipulate import checkpoint_seq, named_apply
  from ._registry import register_model, generate_default_cfgs, register_model_deprecations
label_desc: Use image node_id 1 for calling a global function with example usage: blockify(x, block_size) and returns: x