如何显式广播标量添加维度？类似s[..., np.newaxis]操作

当然可以，numpy里的标量本质是0维ndarray，完全支持用np.newaxis（或等价的None）来显式扩展维度，不用区分标量和数组场景，刚好能帮你去掉分支判断，简化向量化代码。

具体实现示例 #

假设你原来的randomu函数是用分支处理单/多向量场景：

import numpy as np

def randomu(N, M=None):
    if M is None:
        vec = np.random.randn(N)
        norm = np.linalg.norm(vec)
        return vec / norm
    else:
        vecs = np.random.randn(M, N)
        norms = np.linalg.norm(vecs, axis=1)
        return vecs / norms[:, np.newaxis]

现在可以改成无分支的统一写法，完全符合你要的行为：

import numpy as np

def randomu(N, M=None):
    # 生成对应形状的随机数组：M为None时是(N,)，否则是(M,N)
    shape = (N,) if M is None else (M, N)
    vecs = np.random.randn(*shape)
    
    # 计算范数：M为None时返回标量，否则返回(M,)的一维数组
    norms = np.linalg.norm(vecs, axis=-1)
    
    # 统一扩展最后一维：标量变(1,)，数组变(M,1)，自动适配广播
    return vecs / norms[..., np.newaxis]

为什么这能工作？ #

• numpy标量是ndarray的特殊实例（ndim=0），支持所有数组索引操作，scalar[..., np.newaxis]会把它转成形状为(1,)的一维数组。
• 当处理单向量场景时，vecs是(N,)，norms[..., np.newaxis]是(1,)，两者做除法时numpy会自动广播，结果保持(N,)的形状。
• 当处理多向量场景时，vecs是(M,N)，norms[..., np.newaxis]是(M,1)，广播后刚好和vecs的形状匹配，除法结果为(M,N)。

替代写法：np.expand_dims #

如果你觉得索引写法不够直观，也可以用np.expand_dims实现同样效果：

return vecs / np.expand_dims(norms, axis=-1)

这样改写后，你的函数不仅去掉了冗余的分支判断，返回形状也完全和np.zeros、np.random.normal等numpy内置函数对齐，完美支持单/多数据点的统一处理。

内容的提问来源于stack exchange，提问作者LordGössl