tensorflow或pytorch中的Mat2cell matlab等效项

import tensorflow as tf

def mat2cell(matrix, rows, cols):
    # 获取矩阵的形状
    shape = tf.shape(matrix)
    # 计算分块后的子矩阵的形状
    sub_shape = [shape[0] // rows, shape[1] // cols]
    # 使用tf.split函数进行分块操作
    sub_matrices = tf.split(matrix, rows, axis=0)
    sub_matrices = [tf.split(sub_matrix, cols, axis=1) for sub_matrix in sub_matrices]
    # 将子矩阵转换为张量列表
    sub_matrices = [tf.concat(sub_matrix, axis=1) for sub_matrix in sub_matrices]
    return sub_matrices

# 示例用法
matrix = tf.constant([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
sub_matrices = mat2cell(matrix, 2, 2)
print(sub_matrices)

import torch

def mat2cell(matrix, rows, cols):
    # 获取矩阵的形状
    shape = matrix.size()
    # 计算分块后的子矩阵的形状
    sub_shape = [shape[0] // rows, shape[1] // cols]
    # 使用torch.chunk函数进行分块操作
    sub_matrices = torch.chunk(matrix, rows, dim=0)
    sub_matrices = [torch.chunk(sub_matrix, cols, dim=1) for sub_matrix in sub_matrices]
    # 将子矩阵转换为张量列表
    sub_matrices = [torch.cat(sub_matrix, dim=1) for sub_matrix in sub_matrices]
    return sub_matrices

# 示例用法
matrix = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
sub_matrices = mat2cell(matrix, 2, 2)
print(sub_matrices)