Improve documentation of make_causal_mask

The current `make_causal_mask` API is overly restrictive because it requires one to have an array of shape `[B, *, L]`, but `MultiHeadAttention` requires `[B, *, L, H]` as input. This causes a great amount of confusion (as the examples do not have shape annotation/comments) when created mask contains features in its shape and cause shape mismatch errors down the road.

I propose to change the `make_causal_mask` to take a shape tuple as input, instead of an example input array.

As a side issue, many functions have docstrings, but they are not in the generated documents. Most notably are all `__call__` methods, which I think are very important. I am unfamiliar with Sphinx, so I am not sure what kind of issue this would be.