Full pass to strengthen public docstrings in transformer models.
Full pass to strengthen public docstrings in transformer models.

For each model variant:
  - Give TL;DR of what the model does.
  - Characterize model input(s): intended interpretation/use and shapes; optionally further
    restrictions or semantic info.
  - Characterize model output: intended interpretation/use and shape.
  - Repeat arg descriptions across all models/functions that use the same args.

Also:
  - Move max_len keyword arg above dropout-related args; it didn't feel right
    mixed in between dropout settings and mode.
PiperOrigin-RevId: 325955409

google/trax,google/trax