Numpy：对二维数组求带前置条件的点积（求和前应用函数g）

这是标准矩阵乘法的变体，核心逻辑是对矩阵a(n×d)和b(d×h)，结果矩阵c(n×h)的每个元素c[i][k]满足：

c[i][k] = Σ(g(a[i][j] * b[j][k])) ，其中j从0到d-1

以下是几种不同场景下的可行实现方案：

1. 纯Python实现（小矩阵场景） #

适合矩阵规模较小、不需要高性能的场景，逻辑直观：

def custom_matmul_py(a, b, g):
    n = len(a)
    d = len(a[0])
    h = len(b[0])
    # 初始化结果矩阵
    c = [[0.0 for _ in range(h)] for _ in range(n)]
    for i in range(n):
        for k in range(h):
            total = 0
            for j in range(d):
                total += g(a[i][j] * b[j][k])
            c[i][k] = total
    return c

# 测试示例
a = [[1,2,3],[4,5,6]]
b = [[1,2,3],[4,5,6],[1,1,1]]
# 示例g函数：直接返回输入（等价于普通矩阵乘法）
g = lambda x: x
result = custom_matmul_py(a, b, g)
print(result)
# 输出：[[12, 15, 18], [30, 39, 48]]

2. NumPy向量化实现（大数据量高效计算） #

利用NumPy的广播机制实现无循环计算，性能远高于纯Python循环，适合大规模矩阵：

import numpy as np

def custom_matmul_np(a, b, g):
    # 将输入转为NumPy数组
    a_np = np.asarray(a)
    b_np = np.asarray(b)
    # 广播得到(n,d,h)的元素积矩阵
    products = np.expand_dims(a_np, axis=2) * np.expand_dims(b_np, axis=0)
    # 对每个元素应用g，再沿d轴求和
    return np.sum(g(products), axis=1)

# 测试示例
a_np = np.array([[1,2,3],[4,5,6]])
b_np = np.array([[1,2,3],[4,5,6],[1,1,1]])
g = lambda x: x  # 等价普通矩阵乘法
result_np = custom_matmul_np(a_np, b_np, g)
print(result_np)
# 输出：
# [[12 15 18]
#  [30 39 48]]

注意 #

如果g是无法直接处理数组的自定义函数，可先用np.vectorize包装（会损失部分性能）：

def custom_g(x):
    # 示例：自定义非向量化函数
    return x ** 2 if x > 5 else x

g_vectorized = np.vectorize(custom_g)
result = custom_matmul_np(a_np, b_np, g_vectorized)

3. 深度学习框架实现（GPU加速/可微分场景） #

如果需要GPU加速或支持反向传播（比如机器学习任务），可以用PyTorch或TensorFlow实现：

PyTorch版本 #

import torch

def custom_matmul_torch(a, b, g):
    a_torch = torch.tensor(a, dtype=torch.float32)
    b_torch = torch.tensor(b, dtype=torch.float32)
    # 广播得到(n,d,h)的元素积张量
    products = a_torch.unsqueeze(2) * b_torch.unsqueeze(0)
    # 应用g后沿d轴求和
    return torch.sum(g(products), dim=1)

# 测试示例
result_torch = custom_matmul_torch(a, b, lambda x: x)
print(result_torch.numpy())
# 输出与NumPy版本一致

TensorFlow版本 #

import tensorflow as tf

def custom_matmul_tf(a, b, g):
    a_tf = tf.convert_to_tensor(a, dtype=tf.float32)
    b_tf = tf.convert_to_tensor(b, dtype=tf.float32)
    products = tf.expand_dims(a_tf, axis=2) * tf.expand_dims(b_tf, axis=0)
    return tf.reduce_sum(g(products), axis=1)

# 测试示例
result_tf = custom_matmul_tf(a, b, lambda x: x)
print(result_tf.numpy())

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