JAX中实现“移位”矩阵乘法的高效方法

你需要计算的是针对每个t的双重求和：
[s_t=\sum_{\theta=\theta_0}^{{\theta_{N-1}}\sum_{a=0}}{J-1}f_{\theta,t,a}g_{t-a,a}]
其中t-a<0时g取0。原方法通过展开所有索引实现，不仅内存开销大，也没利用JAX的向量化优化能力。下面提供两种更高效优雅的实现方式：

方法一：广播+掩码（推荐） #

核心思路是先对θ维度求和减少计算量，再通过广播生成t和a的索引矩阵，结合掩码处理t-a<0的边界情况，最后完成逐元素相乘求和。

import jax.numpy as jnp

# 假设f (N,T,J)、g (T,J)为输入数组
f_sum = f.sum(axis=0)  # 先对θ维度求和，得到(T,J)形状的数组

# 生成t和a的索引矩阵，通过广播对齐维度
t_indices = jnp.arange(T)[:, None]  # shape (T, 1)
a_indices = jnp.arange(J)[None, :]  # shape (1, J)
t_minus_a = t_indices - a_indices   # shape (T, J)

# 获取合法的g值：t-a≥0时取g[t-a,a]，否则取0
g_valid = jnp.where(t_minus_a >= 0, g[t_minus_a, a_indices], 0)

# 计算最终结果s，shape为(T,)
s = (f_sum * g_valid).sum(axis=1)

这种方式完全基于向量化操作，JAX会自动将其编译为高效的GPU/CPU指令，避免了显式索引展开带来的内存浪费，代码也更简洁易读。

方法二：利用一维卷积特性 #

观察求和式的结构，对于每个a，f_sum[t,a] * g[t-a,a]等价于f_sum[:,a]与g[:,a]在t位置的互相关运算结果。我们可以用jax.numpy.convolve结合vmap批量处理所有a通道：

import jax.numpy as jnp
from jax import vmap

# 先对θ维度求和
f_sum = f.sum(axis=0)  # (T,J)

# 定义单通道的互相关计算：输入f_channel(T,)和g_channel(T,)，返回对应s的贡献
def single_channel_corr(f_channel, g_channel):
    # 互相关等价于卷积前翻转g_channel，这里我们需要的是sum_k f[t]g[t-k]，对应mode='full'后取前T个元素
    corr = jnp.convolve(f_channel, g_channel[::-1], mode='full')[:T]
    return corr

# 对所有J个通道批量计算，再求和得到最终s
channel_contributions = vmap(single_channel_corr)(f_sum.T, g.T)  # (J,T)
s = channel_contributions.sum(axis=0)  # (T,)

这种方法适合熟悉卷积操作的场景，性能和广播法相近，但代码稍复杂一些。

两种方法都比原实现高效得多，尤其是在N、T、J较大的场景下，内存占用和计算速度会有明显提升。

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