Component sizes in vanilla RNN

I would like to seek some clarifications on the dimensionalities of the components and weight parameters in a vanilla RNN model performing text classification for the next word. I will present my understanding which is unclear towards the end. Please point out any part that is wrong.

If each word is vectorised over a vocabulary of size $$K$$, then the input_size = number of features is $$K$$. I have a input $$X$$ of size (batch_size ($$n_x$$), seq_len ($$n$$), input_size($$K$$)). If the hidden layer has dimension $$M$$, then the input-to-hidden weight matrix $$W_{hx}$$ has size ($$M, n_x$$). The hidden-to-hidden weight matrix has size ($$M, M$$).

At step $$t$$, $$x_t$$ has size ($$n_x, K$$). \begin{align} a_t&= W_{hx}x_t + W_{hh}h_{t-1} + b_h \\ h_t &= \text{tanh}(a_t) \end{align} So both the activation $$a_t$$ and the hidden state $$h_t$$ have size ($$M, K$$).

I find it not very clear from this point on. If the hidden-to-output weight matrix $$W_{yh}$$ has size ($$K,M$$), \begin{align} \hat{y}_t &= \text{softmax}(W_{yh}h_t) \end{align} the predicted output at step $$t$$, $$\hat{y}_t$$ by my intuition will have size ($$K,1$$), because each entry is supposed to represent the probability of the $$k$$-th word in the vocabulary appearing next. But here it seems to become ($$K,K$$).

Your input-to-hidden matrix $$W_{hx}$$ has shape $$M \times K$$. Your hidden matrix $$W_{hh}$$ has shape $$M \times M$$. Then $$h_t, b_h, a_t$$ all have shape $$M$$. The output matrix $$W_{yh}$$ has shape $$K \times M$$, so $$W_{yh}h_t$$ has shape $$K$$. Softmax doesn't change any shapes, so your output is $$K$$.
You seemed confused about whether to think about $$x_t, h_t$$, etc as a single input/state or an entire batch of them. It doesn't really matter, as long as you're consistent about it. Above, I wrote it out as if $$x_t, h_t$$, etc are a single input. In batch form, we would have $$x_t$$ is $$B \times K$$, and $$h_t, a_t$$ are $$B \times M$$. Then $$\hat y_t$$ is $$B \times K$$.
Of course in batch form, the notation for matrix multiplication is a bit more convoluted -- it doesn't make sense to write $$W_{hx} x_t$$ anymore, you have to write this as $$x_t W_{hx}^T$$ which now has the correct shape $$B \times M$$.
If you don't want to think about time index $$t$$ either, you could write it out as: $$X$$ is $$B \times N \times K$$, and $$H, A$$ are both $$B \times N \times M$$, and $$\hat Y$$ is $$B \times N \times K$$.