Fix docs for TransformerBlock and remove incompatible functions

[mashu]

• Added examples how to use TransformerBlock (not just how to use constructor)
• Removed confusing cross_att_padding_mask and self_att_padding_mask which were returning in log-space and couldn't be used with apply_pad_mask. Now mask can be passed directly.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 40.43%. Comparing base (bfd2207) to head (1b50bba).

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #37      +/-   ##
==========================================
+ Coverage   31.72%   40.43%   +8.70%     
==========================================
  Files          36       33       -3     
  Lines         769      601     -168     
==========================================
- Hits          244      243       -1     
+ Misses        525      358     -167     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[murrellb]

Note: self_att_padding_mask and cross_att_padding_mask are used with the IPA layers, which (currently) need the full pairwise mask to be materialized, as it gets added into the attention logits prior to softmax:

We can't remove these, but we could maybe improve the docstrings?

[mashu]

@murrellb I restored the padding masks, there is no diff in the code except for import being one-liner, other changes are improved some docstrings, testing that examples actually work (previous ones were wrong).