Hi Bin!



 111414:

the memory complexity of this method are much higher than the trivial matrix product (O(mnp) vs. O(mp)).


Yes, I believe that you are correct about this.
When the tensors are broadcast, I do not believe (to the best of my
knowledge) that the broadcast tensors are actually materialized.
But (again to the best of my knowledge) the “outer-difference” tensor
that then becomes the argument to the .norm() most likely is
materialized, so I do expect the memory cost you describe.
(If you can afford the memory cost, I do expect the broadcast version
to run much faster than a for-loop implementation.)
As an aside, pytorch’s JIT talks about “fusing kernels,” but it talks
specifically about fusing “pointwise” operations, so I don’t think that
JIT applies to your use case.  (I could be wrong.)
Best.
K. Frank