Implementation using einsum
Implementation using einsum(with simultaneous broadcasting of linear functions)
# Local memory contains,
# Q: y k h # K: x k h
X
=
einops
.
einsum
(
Q
,
K
,
'y
k h, x k h
-
>
y x h
'
)
X
=
X
/
math
.
sqrt
(k)
# Local memory contains,
# Q: y k h # K: x k h
# Transpose K,
Q
,
K
=
Q
,
einops
.
einsum
(
K
,
'
x k
h
-
>
k x
h'
)
# Implicit outer product and diagonalize,
X
=
einops
.
einsum
(
Q
,
K
,
'y k
1
h
, k
2
x
h
\
-
> y k1 k2 x
h
'
)
# Inner product,
X
=
einops
.
einsum
(
X
,
'y
k k
x h
-
> y x h'
)
# Scale,
X
=
X
/
math
.
sqrt
(k)