def im2col(context, node):
    """
    Extract sliding local blocks from a batched input tensor (rank=4).

    torch.nn.functional.unfold aims to be the general version: im2col is the rank=4 case of unfold.
    PyTorch currently only supports rank=4 input: torch.nn.functional.unfold redispatches to at::im2col,
    which is why coremltools needs im2col to convert torch.nn.functional.unfold.

    We currently only support rank=4 input (consistent with PyTorch),
    by gathering image entries into columns. More flexbible dilation support will be added in the future.

    Reference https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html
    """

    def _parse_positional_args(context, node) -> Tuple[Var]:
        inputs = _get_inputs(context, node, expected=5)

        x = inputs[0]
        kernel_size = inputs[1]
        dilation = inputs[2]
        padding = inputs[3]
        stride = inputs[4]

        if x.rank != 4:
            raise ValueError("Only supports rank=4 input data for im2col (unfold).")
        return x, kernel_size, dilation, padding, stride

    x, kernel_size, dilation, padding, stride = _parse_positional_args(context, node)
    if isinstance(kernel_size, Var):
        kernel_size = kernel_size.val
    if isinstance(dilation, Var):
        dilation = dilation.val
    if isinstance(padding, Var):
        padding = padding.val
    if isinstance(stride, Var):
        stride = stride.val
    N, C, _, _ = x.shape

    def _create_conv2d_weight(
        in_channels: int,
        kernel_size: tuple[int, int],
        dtype: torch.dtype,
    ) -> np.ndarray:
        size = np.prod(kernel_size).item()
        out_channels = in_channels * size
        weight = torch.zeros((out_channels, in_channels, *kernel_size), dtype=dtype)
        ic_idx, h_idx, w_idx = np.meshgrid(
            np.arange(in_channels),
            np.arange(kernel_size[0]),
            np.arange(kernel_size[1]),
            indexing='ij',
        )
        oc_idx = ic_idx * size + h_idx * kernel_size[1] + w_idx
        weight[oc_idx, ic_idx, h_idx, w_idx] = 1
        return weight.numpy()

    w = _create_conv2d_weight(
        in_channels=C,
        kernel_size=tuple(kernel_size.tolist()),
        dtype=MIL_DTYPE_TO_TORCH_DTYPE[x.dtype],
    )
    conv2d = mb.conv(
        x=x,
        weight=w,
        strides=stride,
        pad_type="custom",
        pad=(padding[0], padding[0], padding[1], padding[1]),
        dilations=dilation,
    )
    output = mb.reshape(x=conv2d, shape=(N, w.shape[0], -1), name=node.name)
    context.add(output)

    Extract sliding local blocks from a batched input tensor (rank=4).

    torch.nn.functional.unfold aims to be the general version: im2col is the rank=4 case of unfold.
    PyTorch currently only supports rank=4 input: torch.nn.functional.unfold redispatches to at::im2col,
    which is why coremltools needs im2col to convert torch.nn.functional.unfold.

    We currently only support rank=4 input (consistent with PyTorch),
    by gathering image entries into columns. More flexbible dilation support will be added in the future.

    Reference https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html
    
coremltools/converters/mil/frontend/torch/ops.py