Why does this multiplication of $Q$ and $K$ have a variance of $d_k$, in scaled dot product attention?
In statistics, if $X$ and $Y$ are independent and randomly distributed variables:
$\mathbb{E}[X + Y] = \mathbb{E}[X] + \mathbb{E}[Y] \\
Var(X + Y) = Var(X) + Var(Y) \\
\mathbb{E}[XY] = \mathbb{E}[X]\mathbb{E}[Y] \\
Var(XY) = (Var(X) + \mathbb{E}[X]^2)(Var(Y) + \mathbb{E}[Y]^2) - \mathbb{E}[X]^2\mathbb{E}[Y]^2$
Let $Q$ and $K$ be random $d_k$ x $d_k$ matrices, where each entry is a some random distribution with $0$ mean and $1$ variance. Every entry is independent from each other.
Since each entry of $Q$ and $K$ have identical distribution, we can focus only on the top-left-most element of $QK$ without loss of generality. The same applies to every other element.
The top-left-most element of $QK$ is $\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1}$.
Since $Q$ and $K$ are independent:
$\mathbb{E}[Q_{1, i} K_{i, 1}] = \mathbb{E}[Q_{1, i}] \mathbb{E}[K_{i, 1}] = 0 \\
Var(Q_{1, i} K_{i, 1}) = (Var(Q_{1, i}) + \mathbb{E}[Q_{1, i}]^2)(Var(K_{i, 1}) + \mathbb{E}[K_{i, 1}]^2) - \mathbb{E}[Q_{1, i}]^2\mathbb{E}[K_{i, 1}]^2 = 1$
And so summing up $d_k$ of them:
$\mathbb{E} \left[\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1} \right] = \displaystyle \sum_{i=0}^{d_k} \mathbb{E} \left[ Q_{1,i} K_{i, 1}  \right] = 0 \\
Var\left(\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1} \right) = \displaystyle \sum_{i=0}^{d_k} Var\left( Q_{1,i} K_{i, 1}  \right) = d_k$
For your code block, you are computing the dot product of matrices $a$ and $b$, when you should be doing a matrix multiplication (the attention function multiplies $Q$ by $K$ after all, which is the vectorized form of dot-product -- it doesn't actually do dot product). It should work out to unit variance.
edit: the last paragraph is incorrect, as dot is the same as matrix multiply in the above case