查询Matlab的shiftdim函数对应的Python等效替代函数

你可以直接使用NumPy库的原生方法实现完全一致的功能，无需额外安装第三方库：

功能匹配逻辑 #

• 当n为正整数（向左偏移n个维度，将前n个维度移动到维度列表末尾）：
  调用np.moveaxis(A, range(n), range(-n, 0))即可，和shiftdim(A,n)效果完全一致
• 当n为负整数（向右偏移n个维度，在数组最前方插入|n|个长度为1的维度）：
  调用np.expand_dims(A, axis=tuple(range(-n)))即可匹配效果
• 当n=0时直接返回原数组即可

可直接复用的封装函数 #

你可以直接使用以下封装好的函数，替换Matlab代码中的shiftdim调用：

import numpy as np

def shiftdim(A, n):
    if n > 0:
        return np.moveaxis(A, list(range(n)), list(range(-n, 0)))
    elif n < 0:
        return np.expand_dims(A, axis=tuple(range(-n)))
    return A

示例验证 #

你给出的示例对应运行效果如下：

# 构造形状为2×3×4的数组
A = np.random.rand(2, 3, 4)
B = shiftdim(A, 2)
print(B.shape) 
# 输出：(4, 2, 3)，和Matlab的`shiftdim(A,2)`返回结果完全一致

如果你使用的是PyTorch张量，替换np.moveaxis为torch.movedim、np.expand_dims为torch.unsqueeze即可，逻辑完全一致。

内容的提问来源于stack exchange，提问作者user15854338