Skip to content

Instantly share code, notes, and snippets.

@qpwo

qpwo/as_strided_matmul.py

Created August 15, 2022 02:04
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save qpwo/a39bd54eb317d442ef18b855db25793b to your computer and use it in GitHub Desktop.
Save qpwo/a39bd54eb317d442ef18b855db25793b to your computer and use it in GitHub Desktop.
Download ZIP
as_strided matrix multiplication
Raw
as_strided_matmul.py
import torch as t
l = 2
n = 3
m = 5
a = t.randint(0,100,(l,n))
b = t.randint(0,100,(n,m))
print("matmul:", a.matmul(b))
print("with strides:", (a.as_strided((l,n,m), (n,1,0)) * b.as_strided((l, n, m), (0, m, 1))).sum(1))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment