Reorganize Reformer2 code in preparation for deeper refactoring.
Reorganize Reformer2 code in preparation for deeper refactoring.

Details:
  - Use inner functions to visually separate setting of multiple option values
    from higher-level layer composition. (Use local variables for layer
    instances only when necessary, e.g., when sharing weights or when
    inspecting layer properties, such as n_out.)
  - Align code between DecoderBlock and EncoderBlock to better highlight
    similarities.
  - Use camel-case names for functions that return layer instances, and
    snake-case names otherwise.
PiperOrigin-RevId: 372668894

google/trax,google/trax