How to achieve efficient Generalized Matrix Product in PyTorch?

111414: <p>I have two real-valued matrices A and B with the shapes of (m, n) and (n, p), respectively.<br>
I know there are fast APIs (e.g. mm, matmul, einsum) in PyTorch to achieve matrix multiplication, but I want to know <strong>if there exists an efficient way to change the dimension-reduction function from inner product to the operation I indicated (e.g. L2 norm of residual)</strong>? <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=10" title=":thinking:" class="emoji" alt=":thinking:"><br>
Specifically, I want to achieve the following function f efficiently:</p>
<pre><code class="lang-auto">import torch

m, n, p = 3, 4, 5

A = torch.rand(m, n)
B = torch.rand(n, p)

def g(a, b):
    return (a - b).norm()

# The mysterious efficient implementation I want
C = f(A, B, g)
</code></pre>
<p>where C is with a shape of (m, p), and the element C[i,j] equals to g(A[i], B[:, j]), and actually g could be with other functions (e.g. trivial summation, L1 norm of concatenated vector).<br>
I made some efforts but my solutions are all with inefficient nested <strong>for</strong>-loops. <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=10" title=":sweat_smile:" class="emoji" alt=":sweat_smile:"></p>
<pre><code class="lang-auto"># My inefficient implementation :(
C = torch.empty(m, p)
for i in range(m):
    for j in range(p):
        C[i, j] = g(A[i], B[:, j])
</code></pre>
<p>I believe developing extra CUDA operators can achieve this, but I am totally ignorant for CUDA-programming and prefer an elegant implementation.<br>
Thanks.</p>

KFrank: <p>Hi Bin!</p>
<aside class="quote no-group quote-modified" data-username="111414" data-post="1" data-topic="138523" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/111414/48/31166_2.png" class="avatar"> 111414:</div>
<blockquote>
<p><strong>if there exists an efficient way to change the dimension-reduction function from inner product to the operation I indicated (e.g. L2 norm of residual)</strong>?</p>
<pre data-code-wrap="Python"><code class="lang-Python">def g(a, b):
    return (a - b).norm()
</code></pre>
</blockquote>
</aside>
<p>I think, in general, if your <code>g()</code> is built from a <em>broadcastable</em> operation<br>
and a <em>reduction</em> operation, you can accomplish your goal with pytorch<br>
broadcasting.</p>
<aside class="quote no-group">
<blockquote>
<p>I made some efforts but my solutions are all with inefficient nested <strong>for</strong>-loops.</p>
</blockquote>
</aside>
<p>In your specific case, you can use <code>unsqueeze()</code> in the appropriate<br>
locations to replace the nested loops with  broadcasting, and then<br>
apply <code>norm()</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt;
&gt;&gt;&gt; m, n, p = 3, 4, 5
&gt;&gt;&gt;
&gt;&gt;&gt; A = torch.rand(m, n)
&gt;&gt;&gt; B = torch.rand(n, p)
&gt;&gt;&gt;
&gt;&gt;&gt; def g(a, b):
...     return (a - b).norm()
...
&gt;&gt;&gt; C = torch.empty(m, p)
&gt;&gt;&gt; for i in range(m):
...     for j in range(p):
...         C[i, j] = g(A[i], B[:, j])
...
&gt;&gt;&gt; C
tensor([[1.0094, 0.6486, 0.8481, 0.7942, 0.4473],
        [1.3406, 1.1867, 0.9334, 0.7280, 0.8513],
        [0.5945, 0.7856, 0.9520, 1.2410, 0.7428]])
&gt;&gt;&gt;
&gt;&gt;&gt; (A.unsqueeze (-1) - B.unsqueeze (0)).norm (dim = 1)
tensor([[1.0094, 0.6486, 0.8481, 0.7942, 0.4473],
        [1.3406, 1.1867, 0.9334, 0.7280, 0.8513],
        [0.5945, 0.7856, 0.9520, 1.2410, 0.7428]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

111414: <p>I think this approach can achieve the function I wanted, but I think when m, n and p are large, the memory complexity of this method are much higher than the trivial matrix product (O(mnp) vs. O(mp)). When this operation is applied into some neural networks (e.g. Transformer), the memory cost may be very huge.</p>

KFrank: <p>Hi Bin!</p>
<aside class="quote no-group" data-username="111414" data-post="3" data-topic="138523" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/111414/48/31166_2.png" class="avatar"> 111414:</div>
<blockquote>
<p>the memory complexity of this method are much higher than the trivial matrix product (O(mnp) vs. O(mp)).</p>
</blockquote>
</aside>
<p>Yes, I believe that you are correct about this.</p>
<p>When the tensors are broadcast, I do not believe (to the best of my<br>
knowledge) that the broadcast tensors are actually materialized.</p>
<p>But (again to the best of my knowledge) the “outer-difference” tensor<br>
that then becomes the argument to the <code>.norm()</code> most likely is<br>
materialized, so I do expect the memory cost you describe.</p>
<p>(If you can afford the memory cost, I do expect the broadcast version<br>
to run much faster than a for-loop implementation.)</p>
<p>As an aside, pytorch’s <a href="https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#fuse-pointwise-operations" rel="noopener nofollow ugc">JIT</a> talks about “fusing kernels,” but it talks<br>
specifically about fusing “pointwise” operations, so I don’t think that<br>
JIT applies to your use case.  (I could be wrong.)</p>
<p>Best.</p>
<p>K. Frank</p>