[doc] Add shape annotation for complex layer

[FrankLeeeee]

For complex layers such as self-attention, we can added shape annoation (e.g. # hidden states (b, s, h)) so that the other maintainer can better understand the model logic.