如何在np.einsum中用省略号实现任意维度张量的逐位矩阵作用？

你的需求是让2×2矩阵X依次作用于任意维度的全2张量（形状为(2,2,...,2)）的每一个维度，目前通过手动写einsum索引字符串实现了3维的情况，但想通过省略号简化操作。这里给出几个更通用的解决方案：

核心问题：einsum省略号的局限性 #

np.einsum的省略号...仅代表所有未显式指定的剩余索引，只能放在索引字符串的开头或结尾，无法直接定位中间位置的维度。所以没法直接用省略号写出"作用于第k个中间维度"的表达式，需要结合其他方法间接实现。

方案1：moveaxis + einsum（最直观） #

思路是先把目标维度移到张量的首尾位置，用省略号完成einsum计算后，再把维度移回原位置：

import numpy as np

def apply_matrix_to_kth_dim(X, tensor, k):
    # 将第k个维度移到第一个位置
    moved_tensor = np.moveaxis(tensor, k, 0)
    # 用einsum完成矩阵与该维度的乘法：X的第二维与moved_tensor的第一维求和
    transformed = np.einsum('ab, b...', X, moved_tensor)
    # 将结果的第一维移回原k位置
    return np.moveaxis(transformed, 0, k)

# 示例：4维张量，依次作用每个维度
X = np.array([[1,2],[3,4]])
n = 4
tensor = np.ones((2,) * n)
current_tensor = tensor.copy()

for k in range(n):
    current_tensor = apply_matrix_to_kth_dim(X, current_tensor, k)

方案2：tensordot + 轴调整（更高效） #

np.tensordot可以直接指定要收缩的轴，配合moveaxis调整结果维度位置，比手动写einsum索引更简洁：

def apply_matrix_to_kth_dim(X, tensor, k):
    # 收缩X的第二维和张量的第k维
    contracted = np.tensordot(X, tensor, axes=([1], [k]))
    # 将X的第一维（结果的第0维）移回原k位置
    return np.moveaxis(contracted, 0, k)

# 使用方式和方案1完全一致

方案3：动态生成einsum索引字符串（兼容你的原始思路） #

如果坚持要用纯einsum的方式，可以动态生成索引字符串，避免手动拼接：

def apply_matrix_to_kth_dim(X, tensor, k):
    n = tensor.ndim
    # 生成足够多的唯一索引字符
    indices = [chr(ord('a') + i) for i in range(n + 1)]
    # X的索引：第一个字符为输出索引，第二个为收缩索引
    x_idx = indices[0] + indices[1]
    # 张量的索引：第k个位置设为收缩索引，其余为原索引
    tensor_idx = ''.join([indices[1] if i == k else indices[i+1] for i in range(n)])
    # 输出索引：替换张量的收缩索引为X的输出索引
    output_idx = ''.join([indices[0] if i == k else indices[i+1] for i in range(n)])
    # 构造einsum字符串
    einsum_str = f'{x_idx}, {tensor_idx} -> {output_idx}'
    return np.einsum(einsum_str, X, tensor)

这个方法无需依赖其他函数，但当张量维度超过26时需要扩展索引字符（比如用多字符索引），适用性不如前两种方案。

总结 #

• 优先选择方案1或方案2，无需手动处理索引字符串，支持任意维度（只要内存允许），代码可读性和通用性更强；
• 如果必须用纯einsum，方案3可以实现动态索引生成，但存在维度数量限制。

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