将逐元素矩阵运算转换为矩阵乘法？

Question

为了利用我的张量核心，我需要编写矩阵乘法（如果我错了，请纠正我）。不幸的是，我的解决方案不是自然地采用 matmul 形式。但也许它可以转换成它。

这是我的计算图的一次迭代：

down = scent.roll(1, 0) * obstaclesMask * topWallMask;
up = scent.roll(-1, 0) * obstaclesMask * bottomWallMask;
right = scent.roll(1, 1) * obstaclesMask * leftWallMask;
left = scent.roll(-1, 1) * obstaclesMask * rightWallMask;
scent = scent + ((down - scent) + (up - scent) + (right - scent) + (left - scent)) * diffusionRate;
scent = torch::max(scent, entities);

（这是 libtorch，但 pytorch 看起来非常相似）。所有这些都是 NxN 矩阵，除了 diffusionRate 是一个标量。我正在考虑向下/向上/向右/向左 - 或一些中间矩阵 - 可以堆叠和组合，但我不确定。想法？