如何在Python中创建仅运行时计算元素的Numpy大尺寸张量？

Numpy原生数组是预分配内存并存储所有元素的，没法直接实现惰性求值（仅在访问时计算元素）。要实现你需要的功能，最直接的方式是自定义一个类来模拟张量的行为，仅在请求元素时触发计算逻辑。

基础实现：单个元素按需计算 #

下面是一个极简的自定义类，模拟多维张量的行为，只有当你访问某个索引的元素时才计算对应的值：

class LazyTensor:
    def __init__(self, shape):
        self.shape = shape  # 存储张量的维度信息
    
    def __getitem__(self, indices):
        # 统一索引格式为元组，兼容单维度访问的情况
        if not isinstance(indices, tuple):
            indices = (indices,)
        
        # 检查索引是否越界
        for idx, dim_size in zip(indices, self.shape):
            if not (0 <= idx < dim_size):
                raise IndexError(f"Index {idx} out of bounds for dimension of size {dim_size}")
        
        # 计算元素值，处理除以0的异常
        sum_indices = sum(indices)
        if sum_indices == 0:
            raise ValueError("Cannot compute value: sum of indices is zero (division by zero)")
        return 1 / sum_indices

# 使用示例
# 创建一个10×10×10×10×10的惰性张量
A = LazyTensor((10, 10, 10, 10, 10))
# 仅计算并返回指定索引的元素
print(A[0, 0, 0, 0, 1])  # 输出: 1.0
print(A[1, 2, 3, 4, 5])  # 输出: 0.06666666666666667

这个类初始化时完全不占用内存，不管你设置的形状多大（比如(100,100,100,100,100)），只有当你访问具体元素时才会执行计算。

扩展实现：支持切片批量计算 #

如果需要支持切片操作（比如A[0:2, 0:2, 0, 0, 0]），可以扩展__getitem__方法，批量计算切片范围内的元素并返回Numpy数组：

import numpy as np

class LazyTensor:
    def __init__(self, shape):
        self.shape = shape
    
    def __getitem__(self, indices):
        # 补全索引，比如如果只传3个索引，自动补全后面的维度为0
        while len(indices) < len(self.shape):
            indices += (0,)
        
        # 处理每个维度的索引/切片
        index_arrays = []
        for idx, dim_size in zip(indices, self.shape):
            if isinstance(idx, slice):
                # 将切片转换为索引数组
                start, stop, step = idx.indices(dim_size)
                index_arrays.append(np.arange(start, stop, step))
            else:
                # 单个索引转为长度为1的数组
                index_arrays.append(np.array([idx]))
        
        # 生成多维网格坐标，对应切片内的所有元素位置
        coords = np.meshgrid(*index_arrays, indexing='ij')
        sum_coords = sum(coords)
        
        # 计算值，忽略除以0的警告（可根据需求调整）
        with np.errstate(divide='ignore'):
            result = 1 / sum_coords
        
        return result

# 切片使用示例
A = LazyTensor((10, 10, 10, 10, 10))
# 获取前2×2的元素切片
slice_result = A[0:2, 0:2, 0, 0, 0]
print(slice_result)
# 输出:
# [[inf 1. ]
#  [1.  0.5]]

注意事项 #

• 如果需要和Numpy的其他函数完全兼容，还可以实现__array__方法，但调用该方法会触发所有元素的计算，对于超大尺寸的张量，依然会面临内存问题，需谨慎使用。
• 自定义类的性能取决于计算逻辑的复杂度，对于你的场景（仅求和取倒数），计算开销可以忽略不计。

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