如何在Tensorflow中复制PyTorch的nn.functional.unfold函数？

import tensorflow as tf

def unfold(input, kernel_size, stride=1, padding=0):
    # 转换输入的维度，将channels_last改为channels_first
    input = tf.transpose(input, [0, 3, 1, 2])
    
    # 计算填充大小
    pad_size = padding if isinstance(padding, int) else padding[0]
    
    # 使用tf.extract_image_patches函数提取图像块
    patches = tf.extract_image_patches(images=input,
                                       ksizes=[1, kernel_size, kernel_size, 1],
                                       strides=[1, stride, stride, 1],
                                       rates=[1, 1, 1, 1],
                                       padding='VALID' if pad_size == 0 else 'SAME')
    
    # 将提取的图像块转换为正确的维度
    batch_size, in_channels, in_height, in_width = input.get_shape().as_list()
    out_channels = in_channels * kernel_size * kernel_size
    out_height = (in_height - kernel_size + 2 * pad_size) // stride + 1
    out_width = (in_width - kernel_size + 2 * pad_size) // stride + 1
    
    patches = tf.reshape(patches, [-1, out_height, out_width, out_channels])
    patches = tf.transpose(patches, [0, 3, 1, 2])
    
    return patches