批量大于1时如何实现torch.linalg.matmul？解决torch.matrix_power批量适配问题

torch.linalg.matmul（或简写@）本身支持批量输入，只需保证张量维度兼容，无需额外处理，以下是两种常见场景的实现：

• 场景1：单权重矩阵对应批量输入
  若输入X形状为(batch_size, dim_x)，权重矩阵A形状为(dim_x, dim_x)，直接执行矩阵乘法即可，PyTorch会自动对批量维度广播计算：

  output = torch.linalg.matmul(X, A)
  # 或简写为 output = X @ A
  

• 场景2：批量权重矩阵对应批量输入
  若A是批量矩阵（形状(batch_size, dim_x, dim_x)），X为(batch_size, dim_x)，需先给X添加一个维度以匹配矩阵乘法的维度要求，计算后再移除多余维度：

  # 调整X维度为 (batch_size, 1, dim_x)
  X_expanded = X.unsqueeze(1)
  output = torch.linalg.matmul(X_expanded, A).squeeze(1)
  # 或简写为 output = (X.unsqueeze(1) @ A).squeeze(1)
  

针对torch.matrix_power仅支持标量幂次、预计算出错的问题，结合k∈[1,10]的限制，提供两种实用方案：

方案1：按需计算（推荐，小幂次开销可忽略） #

由于k的取值范围极小，torch.matrix_power的计算开销极低，直接在forward中计算即可避免预计算的同步问题。若k是批量张量（每个样本对应不同k值），可通过分组批量处理：

import numpy as np
import torch
import torch.nn as nn

class NN(nn.Module):
    def __init__(self, dim_x):
        super().__init__()
        self.A = nn.Parameter(torch.randn(dim_x, dim_x))

    def forward(self, X, k):
        k_vals = k.cpu().numpy()
        output = torch.zeros_like(X)
        # 遍历所有唯一k值，批量处理对应样本
        for k_val in np.unique(k_vals):
            mask = k_vals == k_val
            A_pow = torch.matrix_power(self.A, k_val)
            output[mask] = X[mask] @ A_pow
        return output

方案2：预计算+自动同步（适合重复调用场景） #

若需预计算减少重复运算，可通过哈希值判断参数A是否更新，自动同步预计算的幂矩阵：

import torch
import torch.nn as nn

class NN(nn.Module):
    def __init__(self, dim_x):
        super().__init__()
        self.A = nn.Parameter(torch.randn(dim_x, dim_x))
        self.max_k = 10
        self._cached_powers = None
        self._last_A_hash = None

    def _update_powers(self):
        # 通过哈希值检测A是否更新
        current_hash = hash(self.A.detach().cpu().numpy().tobytes())
        if self._last_A_hash != current_hash:
            self._cached_powers = [torch.matrix_power(self.A, k) for k in range(1, self.max_k+1)]
            self._last_A_hash = current_hash

    def forward(self, X, k):
        self._update_powers()
        # 处理k为张量或标量的情况
        k_val = k.item() if isinstance(k, torch.Tensor) else k
        A_pow = self._cached_powers[k_val - 1]
        return X @ A_pow

注意：若k是全局标量，方案2可有效减少重复计算；若k是批量张量，方案1的分组处理更适配。

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