numpy为数组添加可变维数

Question

有没有办法在numpy数组中添加可变数量的单例维度？我想要atleast_2d之类的东西，但需要任意数量的维度。 expand_dims出现了很多，但只增加了一个维度。

我知道这样做的唯一方法是首先计算形状并应用它，即

import numpy as np

def atleast_kd(array, k):
    array = np.asarray(array)
    new_shape = array.shape + (1,) * (k - array.ndim)
    return array.reshape(new_shape)

Answer 1

更优雅的方法是np.broadcast_to。例如：

a = np.random.rand(2,2)
k = # number_extra dimensions
b = np.broadcast_to(a, (1,) * k + a.shape)

这会将尺寸添加到b的开头。要获得与给定功能完全相同的行为，可以使用np.moveaxis。